jnkr36 commited on
Commit
4bb060d
1 Parent(s): d8bbeea

Upload 71 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. comfy/cldm/cldm.py +283 -0
  2. comfy/clip_vision.py +62 -0
  3. comfy/clip_vision_config_h.json +18 -0
  4. comfy/clip_vision_config_vitl.json +18 -0
  5. comfy/extra_samplers/uni_pc.py +883 -0
  6. comfy/k_diffusion/augmentation.py +105 -0
  7. comfy/k_diffusion/config.py +110 -0
  8. comfy/k_diffusion/evaluation.py +134 -0
  9. comfy/k_diffusion/external.py +179 -0
  10. comfy/k_diffusion/gns.py +99 -0
  11. comfy/k_diffusion/layers.py +246 -0
  12. comfy/k_diffusion/models/__init__.py +1 -0
  13. comfy/k_diffusion/models/image_v1.py +156 -0
  14. comfy/k_diffusion/sampling.py +607 -0
  15. comfy/k_diffusion/utils.py +332 -0
  16. comfy/ldm/data/__init__.py +0 -0
  17. comfy/ldm/data/util.py +24 -0
  18. comfy/ldm/models/autoencoder.py +223 -0
  19. comfy/ldm/models/diffusion/__init__.py +0 -0
  20. comfy/ldm/models/diffusion/ddim.py +410 -0
  21. comfy/ldm/models/diffusion/ddpm.py +1875 -0
  22. comfy/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  23. comfy/ldm/models/diffusion/dpm_solver/dpm_solver.py +1163 -0
  24. comfy/ldm/models/diffusion/dpm_solver/sampler.py +96 -0
  25. comfy/ldm/models/diffusion/plms.py +245 -0
  26. comfy/ldm/models/diffusion/sampling_util.py +22 -0
  27. comfy/ldm/modules/attention.py +591 -0
  28. comfy/ldm/modules/diffusionmodules/__init__.py +0 -0
  29. comfy/ldm/modules/diffusionmodules/model.py +942 -0
  30. comfy/ldm/modules/diffusionmodules/openaimodel.py +821 -0
  31. comfy/ldm/modules/diffusionmodules/upscaling.py +81 -0
  32. comfy/ldm/modules/diffusionmodules/util.py +278 -0
  33. comfy/ldm/modules/distributions/__init__.py +0 -0
  34. comfy/ldm/modules/distributions/distributions.py +92 -0
  35. comfy/ldm/modules/ema.py +80 -0
  36. comfy/ldm/modules/encoders/__init__.py +0 -0
  37. comfy/ldm/modules/encoders/kornia_functions.py +59 -0
  38. comfy/ldm/modules/encoders/modules.py +314 -0
  39. comfy/ldm/modules/encoders/noise_aug_modules.py +35 -0
  40. comfy/ldm/modules/image_degradation/__init__.py +2 -0
  41. comfy/ldm/modules/image_degradation/bsrgan.py +730 -0
  42. comfy/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  43. comfy/ldm/modules/image_degradation/utils/test.png +0 -0
  44. comfy/ldm/modules/image_degradation/utils_image.py +916 -0
  45. comfy/ldm/modules/midas/__init__.py +0 -0
  46. comfy/ldm/modules/midas/api.py +170 -0
  47. comfy/ldm/modules/midas/midas/__init__.py +0 -0
  48. comfy/ldm/modules/midas/midas/base_model.py +16 -0
  49. comfy/ldm/modules/midas/midas/blocks.py +342 -0
  50. comfy/ldm/modules/midas/midas/dpt_depth.py +109 -0
comfy/cldm/cldm.py ADDED
@@ -0,0 +1,283 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #taken from: https://github.com/lllyasviel/ControlNet
2
+ #and modified
3
+
4
+ import torch
5
+ import torch as th
6
+ import torch.nn as nn
7
+
8
+ from ldm.modules.diffusionmodules.util import (
9
+ conv_nd,
10
+ linear,
11
+ zero_module,
12
+ timestep_embedding,
13
+ )
14
+
15
+ from ldm.modules.attention import SpatialTransformer
16
+ from ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
17
+ from ldm.models.diffusion.ddpm import LatentDiffusion
18
+ from ldm.util import log_txt_as_img, exists, instantiate_from_config
19
+
20
+
21
+ class ControlledUnetModel(UNetModel):
22
+ #implemented in the ldm unet
23
+ pass
24
+
25
+ class ControlNet(nn.Module):
26
+ def __init__(
27
+ self,
28
+ image_size,
29
+ in_channels,
30
+ model_channels,
31
+ hint_channels,
32
+ num_res_blocks,
33
+ attention_resolutions,
34
+ dropout=0,
35
+ channel_mult=(1, 2, 4, 8),
36
+ conv_resample=True,
37
+ dims=2,
38
+ use_checkpoint=False,
39
+ use_fp16=False,
40
+ num_heads=-1,
41
+ num_head_channels=-1,
42
+ num_heads_upsample=-1,
43
+ use_scale_shift_norm=False,
44
+ resblock_updown=False,
45
+ use_new_attention_order=False,
46
+ use_spatial_transformer=False, # custom transformer support
47
+ transformer_depth=1, # custom transformer support
48
+ context_dim=None, # custom transformer support
49
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
50
+ legacy=True,
51
+ disable_self_attentions=None,
52
+ num_attention_blocks=None,
53
+ disable_middle_self_attn=False,
54
+ use_linear_in_transformer=False,
55
+ ):
56
+ super().__init__()
57
+ if use_spatial_transformer:
58
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
59
+
60
+ if context_dim is not None:
61
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
62
+ # from omegaconf.listconfig import ListConfig
63
+ # if type(context_dim) == ListConfig:
64
+ # context_dim = list(context_dim)
65
+
66
+ if num_heads_upsample == -1:
67
+ num_heads_upsample = num_heads
68
+
69
+ if num_heads == -1:
70
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
71
+
72
+ if num_head_channels == -1:
73
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
74
+
75
+ self.dims = dims
76
+ self.image_size = image_size
77
+ self.in_channels = in_channels
78
+ self.model_channels = model_channels
79
+ if isinstance(num_res_blocks, int):
80
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
81
+ else:
82
+ if len(num_res_blocks) != len(channel_mult):
83
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
84
+ "as a list/tuple (per-level) with the same length as channel_mult")
85
+ self.num_res_blocks = num_res_blocks
86
+ if disable_self_attentions is not None:
87
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
88
+ assert len(disable_self_attentions) == len(channel_mult)
89
+ if num_attention_blocks is not None:
90
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
91
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
92
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
93
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
94
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
95
+ f"attention will still not be set.")
96
+
97
+ self.attention_resolutions = attention_resolutions
98
+ self.dropout = dropout
99
+ self.channel_mult = channel_mult
100
+ self.conv_resample = conv_resample
101
+ self.use_checkpoint = use_checkpoint
102
+ self.dtype = th.float16 if use_fp16 else th.float32
103
+ self.num_heads = num_heads
104
+ self.num_head_channels = num_head_channels
105
+ self.num_heads_upsample = num_heads_upsample
106
+ self.predict_codebook_ids = n_embed is not None
107
+
108
+ time_embed_dim = model_channels * 4
109
+ self.time_embed = nn.Sequential(
110
+ linear(model_channels, time_embed_dim),
111
+ nn.SiLU(),
112
+ linear(time_embed_dim, time_embed_dim),
113
+ )
114
+
115
+ self.input_blocks = nn.ModuleList(
116
+ [
117
+ TimestepEmbedSequential(
118
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
119
+ )
120
+ ]
121
+ )
122
+ self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])
123
+
124
+ self.input_hint_block = TimestepEmbedSequential(
125
+ conv_nd(dims, hint_channels, 16, 3, padding=1),
126
+ nn.SiLU(),
127
+ conv_nd(dims, 16, 16, 3, padding=1),
128
+ nn.SiLU(),
129
+ conv_nd(dims, 16, 32, 3, padding=1, stride=2),
130
+ nn.SiLU(),
131
+ conv_nd(dims, 32, 32, 3, padding=1),
132
+ nn.SiLU(),
133
+ conv_nd(dims, 32, 96, 3, padding=1, stride=2),
134
+ nn.SiLU(),
135
+ conv_nd(dims, 96, 96, 3, padding=1),
136
+ nn.SiLU(),
137
+ conv_nd(dims, 96, 256, 3, padding=1, stride=2),
138
+ nn.SiLU(),
139
+ zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
140
+ )
141
+
142
+ self._feature_size = model_channels
143
+ input_block_chans = [model_channels]
144
+ ch = model_channels
145
+ ds = 1
146
+ for level, mult in enumerate(channel_mult):
147
+ for nr in range(self.num_res_blocks[level]):
148
+ layers = [
149
+ ResBlock(
150
+ ch,
151
+ time_embed_dim,
152
+ dropout,
153
+ out_channels=mult * model_channels,
154
+ dims=dims,
155
+ use_checkpoint=use_checkpoint,
156
+ use_scale_shift_norm=use_scale_shift_norm,
157
+ )
158
+ ]
159
+ ch = mult * model_channels
160
+ if ds in attention_resolutions:
161
+ if num_head_channels == -1:
162
+ dim_head = ch // num_heads
163
+ else:
164
+ num_heads = ch // num_head_channels
165
+ dim_head = num_head_channels
166
+ if legacy:
167
+ #num_heads = 1
168
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
169
+ if exists(disable_self_attentions):
170
+ disabled_sa = disable_self_attentions[level]
171
+ else:
172
+ disabled_sa = False
173
+
174
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
175
+ layers.append(
176
+ AttentionBlock(
177
+ ch,
178
+ use_checkpoint=use_checkpoint,
179
+ num_heads=num_heads,
180
+ num_head_channels=dim_head,
181
+ use_new_attention_order=use_new_attention_order,
182
+ ) if not use_spatial_transformer else SpatialTransformer(
183
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
184
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
185
+ use_checkpoint=use_checkpoint
186
+ )
187
+ )
188
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
189
+ self.zero_convs.append(self.make_zero_conv(ch))
190
+ self._feature_size += ch
191
+ input_block_chans.append(ch)
192
+ if level != len(channel_mult) - 1:
193
+ out_ch = ch
194
+ self.input_blocks.append(
195
+ TimestepEmbedSequential(
196
+ ResBlock(
197
+ ch,
198
+ time_embed_dim,
199
+ dropout,
200
+ out_channels=out_ch,
201
+ dims=dims,
202
+ use_checkpoint=use_checkpoint,
203
+ use_scale_shift_norm=use_scale_shift_norm,
204
+ down=True,
205
+ )
206
+ if resblock_updown
207
+ else Downsample(
208
+ ch, conv_resample, dims=dims, out_channels=out_ch
209
+ )
210
+ )
211
+ )
212
+ ch = out_ch
213
+ input_block_chans.append(ch)
214
+ self.zero_convs.append(self.make_zero_conv(ch))
215
+ ds *= 2
216
+ self._feature_size += ch
217
+
218
+ if num_head_channels == -1:
219
+ dim_head = ch // num_heads
220
+ else:
221
+ num_heads = ch // num_head_channels
222
+ dim_head = num_head_channels
223
+ if legacy:
224
+ #num_heads = 1
225
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
226
+ self.middle_block = TimestepEmbedSequential(
227
+ ResBlock(
228
+ ch,
229
+ time_embed_dim,
230
+ dropout,
231
+ dims=dims,
232
+ use_checkpoint=use_checkpoint,
233
+ use_scale_shift_norm=use_scale_shift_norm,
234
+ ),
235
+ AttentionBlock(
236
+ ch,
237
+ use_checkpoint=use_checkpoint,
238
+ num_heads=num_heads,
239
+ num_head_channels=dim_head,
240
+ use_new_attention_order=use_new_attention_order,
241
+ ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
242
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
243
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
244
+ use_checkpoint=use_checkpoint
245
+ ),
246
+ ResBlock(
247
+ ch,
248
+ time_embed_dim,
249
+ dropout,
250
+ dims=dims,
251
+ use_checkpoint=use_checkpoint,
252
+ use_scale_shift_norm=use_scale_shift_norm,
253
+ ),
254
+ )
255
+ self.middle_block_out = self.make_zero_conv(ch)
256
+ self._feature_size += ch
257
+
258
+ def make_zero_conv(self, channels):
259
+ return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))
260
+
261
+ def forward(self, x, hint, timesteps, context, **kwargs):
262
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
263
+ emb = self.time_embed(t_emb)
264
+
265
+ guided_hint = self.input_hint_block(hint, emb, context)
266
+
267
+ outs = []
268
+
269
+ h = x.type(self.dtype)
270
+ for module, zero_conv in zip(self.input_blocks, self.zero_convs):
271
+ if guided_hint is not None:
272
+ h = module(h, emb, context)
273
+ h += guided_hint
274
+ guided_hint = None
275
+ else:
276
+ h = module(h, emb, context)
277
+ outs.append(zero_conv(h, emb, context))
278
+
279
+ h = self.middle_block(h, emb, context)
280
+ outs.append(self.middle_block_out(h, emb, context))
281
+
282
+ return outs
283
+
comfy/clip_vision.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CLIPVisionModelWithProjection, CLIPVisionConfig, CLIPImageProcessor
2
+ from .utils import load_torch_file, transformers_convert
3
+ import os
4
+
5
+ class ClipVisionModel():
6
+ def __init__(self, json_config):
7
+ config = CLIPVisionConfig.from_json_file(json_config)
8
+ self.model = CLIPVisionModelWithProjection(config)
9
+ self.processor = CLIPImageProcessor(crop_size=224,
10
+ do_center_crop=True,
11
+ do_convert_rgb=True,
12
+ do_normalize=True,
13
+ do_resize=True,
14
+ image_mean=[ 0.48145466,0.4578275,0.40821073],
15
+ image_std=[0.26862954,0.26130258,0.27577711],
16
+ resample=3, #bicubic
17
+ size=224)
18
+
19
+ def load_sd(self, sd):
20
+ self.model.load_state_dict(sd, strict=False)
21
+
22
+ def encode_image(self, image):
23
+ inputs = self.processor(images=[image[0]], return_tensors="pt")
24
+ outputs = self.model(**inputs)
25
+ return outputs
26
+
27
+ def convert_to_transformers(sd):
28
+ sd_k = sd.keys()
29
+ if "embedder.model.visual.transformer.resblocks.0.attn.in_proj_weight" in sd_k:
30
+ keys_to_replace = {
31
+ "embedder.model.visual.class_embedding": "vision_model.embeddings.class_embedding",
32
+ "embedder.model.visual.conv1.weight": "vision_model.embeddings.patch_embedding.weight",
33
+ "embedder.model.visual.positional_embedding": "vision_model.embeddings.position_embedding.weight",
34
+ "embedder.model.visual.ln_post.bias": "vision_model.post_layernorm.bias",
35
+ "embedder.model.visual.ln_post.weight": "vision_model.post_layernorm.weight",
36
+ "embedder.model.visual.ln_pre.bias": "vision_model.pre_layrnorm.bias",
37
+ "embedder.model.visual.ln_pre.weight": "vision_model.pre_layrnorm.weight",
38
+ }
39
+
40
+ for x in keys_to_replace:
41
+ if x in sd_k:
42
+ sd[keys_to_replace[x]] = sd.pop(x)
43
+
44
+ if "embedder.model.visual.proj" in sd_k:
45
+ sd['visual_projection.weight'] = sd.pop("embedder.model.visual.proj").transpose(0, 1)
46
+
47
+ sd = transformers_convert(sd, "embedder.model.visual", "vision_model", 32)
48
+ return sd
49
+
50
+ def load_clipvision_from_sd(sd):
51
+ sd = convert_to_transformers(sd)
52
+ if "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
53
+ json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_h.json")
54
+ else:
55
+ json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
56
+ clip = ClipVisionModel(json_config)
57
+ clip.load_sd(sd)
58
+ return clip
59
+
60
+ def load(ckpt_path):
61
+ sd = load_torch_file(ckpt_path)
62
+ return load_clipvision_from_sd(sd)
comfy/clip_vision_config_h.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "attention_dropout": 0.0,
3
+ "dropout": 0.0,
4
+ "hidden_act": "gelu",
5
+ "hidden_size": 1280,
6
+ "image_size": 224,
7
+ "initializer_factor": 1.0,
8
+ "initializer_range": 0.02,
9
+ "intermediate_size": 5120,
10
+ "layer_norm_eps": 1e-05,
11
+ "model_type": "clip_vision_model",
12
+ "num_attention_heads": 16,
13
+ "num_channels": 3,
14
+ "num_hidden_layers": 32,
15
+ "patch_size": 14,
16
+ "projection_dim": 1024,
17
+ "torch_dtype": "float32"
18
+ }
comfy/clip_vision_config_vitl.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "attention_dropout": 0.0,
3
+ "dropout": 0.0,
4
+ "hidden_act": "quick_gelu",
5
+ "hidden_size": 1024,
6
+ "image_size": 224,
7
+ "initializer_factor": 1.0,
8
+ "initializer_range": 0.02,
9
+ "intermediate_size": 4096,
10
+ "layer_norm_eps": 1e-05,
11
+ "model_type": "clip_vision_model",
12
+ "num_attention_heads": 16,
13
+ "num_channels": 3,
14
+ "num_hidden_layers": 24,
15
+ "patch_size": 14,
16
+ "projection_dim": 768,
17
+ "torch_dtype": "float32"
18
+ }
comfy/extra_samplers/uni_pc.py ADDED
@@ -0,0 +1,883 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #code taken from: https://github.com/wl-zhao/UniPC and modified
2
+
3
+ import torch
4
+ import torch.nn.functional as F
5
+ import math
6
+
7
+ from tqdm.auto import trange, tqdm
8
+
9
+
10
+ class NoiseScheduleVP:
11
+ def __init__(
12
+ self,
13
+ schedule='discrete',
14
+ betas=None,
15
+ alphas_cumprod=None,
16
+ continuous_beta_0=0.1,
17
+ continuous_beta_1=20.,
18
+ ):
19
+ """Create a wrapper class for the forward SDE (VP type).
20
+
21
+ ***
22
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
23
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
24
+ ***
25
+
26
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
27
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
28
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
29
+
30
+ log_alpha_t = self.marginal_log_mean_coeff(t)
31
+ sigma_t = self.marginal_std(t)
32
+ lambda_t = self.marginal_lambda(t)
33
+
34
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
35
+
36
+ t = self.inverse_lambda(lambda_t)
37
+
38
+ ===============================================================
39
+
40
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
41
+
42
+ 1. For discrete-time DPMs:
43
+
44
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
45
+ t_i = (i + 1) / N
46
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
47
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
48
+
49
+ Args:
50
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
51
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
52
+
53
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
54
+
55
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
56
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
57
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
58
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
59
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
60
+ and
61
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
62
+
63
+
64
+ 2. For continuous-time DPMs:
65
+
66
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
67
+ schedule are the default settings in DDPM and improved-DDPM:
68
+
69
+ Args:
70
+ beta_min: A `float` number. The smallest beta for the linear schedule.
71
+ beta_max: A `float` number. The largest beta for the linear schedule.
72
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
73
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
74
+ T: A `float` number. The ending time of the forward process.
75
+
76
+ ===============================================================
77
+
78
+ Args:
79
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
80
+ 'linear' or 'cosine' for continuous-time DPMs.
81
+ Returns:
82
+ A wrapper object of the forward SDE (VP type).
83
+
84
+ ===============================================================
85
+
86
+ Example:
87
+
88
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
89
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
90
+
91
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
92
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
93
+
94
+ # For continuous-time DPMs (VPSDE), linear schedule:
95
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
96
+
97
+ """
98
+
99
+ if schedule not in ['discrete', 'linear', 'cosine']:
100
+ raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))
101
+
102
+ self.schedule = schedule
103
+ if schedule == 'discrete':
104
+ if betas is not None:
105
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
106
+ else:
107
+ assert alphas_cumprod is not None
108
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
109
+ self.total_N = len(log_alphas)
110
+ self.T = 1.
111
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
112
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
113
+ else:
114
+ self.total_N = 1000
115
+ self.beta_0 = continuous_beta_0
116
+ self.beta_1 = continuous_beta_1
117
+ self.cosine_s = 0.008
118
+ self.cosine_beta_max = 999.
119
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (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), self.log_alpha_array.to(t.device)).reshape((-1))
135
+ elif self.schedule == 'linear':
136
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
137
+ elif self.schedule == 'cosine':
138
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
139
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
140
+ return log_alpha_t
141
+
142
+ def marginal_alpha(self, t):
143
+ """
144
+ Compute alpha_t of a given continuous-time label t in [0, T].
145
+ """
146
+ return torch.exp(self.marginal_log_mean_coeff(t))
147
+
148
+ def marginal_std(self, t):
149
+ """
150
+ Compute sigma_t of a given continuous-time label t in [0, T].
151
+ """
152
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
153
+
154
+ def marginal_lambda(self, t):
155
+ """
156
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
157
+ """
158
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
159
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
160
+ return log_mean_coeff - log_std
161
+
162
+ def inverse_lambda(self, lamb):
163
+ """
164
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
165
+ """
166
+ if self.schedule == 'linear':
167
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
168
+ Delta = self.beta_0**2 + tmp
169
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
170
+ elif self.schedule == 'discrete':
171
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
172
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
173
+ return t.reshape((-1,))
174
+ else:
175
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
176
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
177
+ t = t_fn(log_alpha)
178
+ return t
179
+
180
+
181
+ def model_wrapper(
182
+ model,
183
+ sampling_function,
184
+ noise_schedule,
185
+ model_type="noise",
186
+ model_kwargs={},
187
+ guidance_type="uncond",
188
+ condition=None,
189
+ unconditional_condition=None,
190
+ guidance_scale=1.,
191
+ classifier_fn=None,
192
+ classifier_kwargs={},
193
+ ):
194
+ """Create a wrapper function for the noise prediction model.
195
+
196
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
197
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
198
+
199
+ We support four types of the diffusion model by setting `model_type`:
200
+
201
+ 1. "noise": noise prediction model. (Trained by predicting noise).
202
+
203
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
204
+
205
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
206
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
207
+
208
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
209
+ arXiv preprint arXiv:2202.00512 (2022).
210
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
211
+ arXiv preprint arXiv:2210.02303 (2022).
212
+
213
+ 4. "score": marginal score function. (Trained by denoising score matching).
214
+ Note that the score function and the noise prediction model follows a simple relationship:
215
+ ```
216
+ noise(x_t, t) = -sigma_t * score(x_t, t)
217
+ ```
218
+
219
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
220
+ 1. "uncond": unconditional sampling by DPMs.
221
+ The input `model` has the following format:
222
+ ``
223
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
224
+ ``
225
+
226
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
227
+ The input `model` has the following format:
228
+ ``
229
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
230
+ ``
231
+
232
+ The input `classifier_fn` has the following format:
233
+ ``
234
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
235
+ ``
236
+
237
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
238
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
239
+
240
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
241
+ The input `model` has the following format:
242
+ ``
243
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
244
+ ``
245
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
246
+
247
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
248
+ arXiv preprint arXiv:2207.12598 (2022).
249
+
250
+
251
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
252
+ or continuous-time labels (i.e. epsilon to T).
253
+
254
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
255
+ ``
256
+ def model_fn(x, t_continuous) -> noise:
257
+ t_input = get_model_input_time(t_continuous)
258
+ return noise_pred(model, x, t_input, **model_kwargs)
259
+ ``
260
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
261
+
262
+ ===============================================================
263
+
264
+ Args:
265
+ model: A diffusion model with the corresponding format described above.
266
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
267
+ model_type: A `str`. The parameterization type of the diffusion model.
268
+ "noise" or "x_start" or "v" or "score".
269
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
270
+ guidance_type: A `str`. The type of the guidance for sampling.
271
+ "uncond" or "classifier" or "classifier-free".
272
+ condition: A pytorch tensor. The condition for the guided sampling.
273
+ Only used for "classifier" or "classifier-free" guidance type.
274
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
275
+ Only used for "classifier-free" guidance type.
276
+ guidance_scale: A `float`. The scale for the guided sampling.
277
+ classifier_fn: A classifier function. Only used for the classifier guidance.
278
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
279
+ Returns:
280
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
281
+ """
282
+
283
+ def get_model_input_time(t_continuous):
284
+ """
285
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
286
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
287
+ For continuous-time DPMs, we just use `t_continuous`.
288
+ """
289
+ if noise_schedule.schedule == 'discrete':
290
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
291
+ else:
292
+ return t_continuous
293
+
294
+ def noise_pred_fn(x, t_continuous, cond=None):
295
+ if t_continuous.reshape((-1,)).shape[0] == 1:
296
+ t_continuous = t_continuous.expand((x.shape[0]))
297
+ t_input = get_model_input_time(t_continuous)
298
+ output = sampling_function(model, x, t_input, **model_kwargs)
299
+ if model_type == "noise":
300
+ return output
301
+ elif model_type == "x_start":
302
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
303
+ dims = x.dim()
304
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
305
+ elif model_type == "v":
306
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
307
+ dims = x.dim()
308
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
309
+ elif model_type == "score":
310
+ sigma_t = noise_schedule.marginal_std(t_continuous)
311
+ dims = x.dim()
312
+ return -expand_dims(sigma_t, dims) * output
313
+
314
+ def cond_grad_fn(x, t_input):
315
+ """
316
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
317
+ """
318
+ with torch.enable_grad():
319
+ x_in = x.detach().requires_grad_(True)
320
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
321
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
322
+
323
+ def model_fn(x, t_continuous):
324
+ """
325
+ The noise predicition model function that is used for DPM-Solver.
326
+ """
327
+ if t_continuous.reshape((-1,)).shape[0] == 1:
328
+ t_continuous = t_continuous.expand((x.shape[0]))
329
+ if guidance_type == "uncond":
330
+ return noise_pred_fn(x, t_continuous)
331
+ elif guidance_type == "classifier":
332
+ assert classifier_fn is not None
333
+ t_input = get_model_input_time(t_continuous)
334
+ cond_grad = cond_grad_fn(x, t_input)
335
+ sigma_t = noise_schedule.marginal_std(t_continuous)
336
+ noise = noise_pred_fn(x, t_continuous)
337
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
338
+ elif guidance_type == "classifier-free":
339
+ if guidance_scale == 1. or unconditional_condition is None:
340
+ return noise_pred_fn(x, t_continuous, cond=condition)
341
+ else:
342
+ x_in = torch.cat([x] * 2)
343
+ t_in = torch.cat([t_continuous] * 2)
344
+ c_in = torch.cat([unconditional_condition, condition])
345
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
346
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
347
+
348
+ assert model_type in ["noise", "x_start", "v"]
349
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
350
+ return model_fn
351
+
352
+
353
+ class UniPC:
354
+ def __init__(
355
+ self,
356
+ model_fn,
357
+ noise_schedule,
358
+ predict_x0=True,
359
+ thresholding=False,
360
+ max_val=1.,
361
+ variant='bh1',
362
+ noise_mask=None,
363
+ masked_image=None,
364
+ noise=None,
365
+ ):
366
+ """Construct a UniPC.
367
+
368
+ We support both data_prediction and noise_prediction.
369
+ """
370
+ self.model = model_fn
371
+ self.noise_schedule = noise_schedule
372
+ self.variant = variant
373
+ self.predict_x0 = predict_x0
374
+ self.thresholding = thresholding
375
+ self.max_val = max_val
376
+ self.noise_mask = noise_mask
377
+ self.masked_image = masked_image
378
+ self.noise = noise
379
+
380
+ def dynamic_thresholding_fn(self, x0, t=None):
381
+ """
382
+ The dynamic thresholding method.
383
+ """
384
+ dims = x0.dim()
385
+ p = self.dynamic_thresholding_ratio
386
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
387
+ s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
388
+ x0 = torch.clamp(x0, -s, s) / s
389
+ return x0
390
+
391
+ def noise_prediction_fn(self, x, t):
392
+ """
393
+ Return the noise prediction model.
394
+ """
395
+ if self.noise_mask is not None:
396
+ return self.model(x, t) * self.noise_mask
397
+ else:
398
+ return self.model(x, t)
399
+
400
+ def data_prediction_fn(self, x, t):
401
+ """
402
+ Return the data prediction model (with thresholding).
403
+ """
404
+ noise = self.noise_prediction_fn(x, t)
405
+ dims = x.dim()
406
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
407
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
408
+ if self.thresholding:
409
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
410
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
411
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
412
+ x0 = torch.clamp(x0, -s, s) / s
413
+ if self.noise_mask is not None:
414
+ x0 = x0 * self.noise_mask + (1. - self.noise_mask) * self.masked_image
415
+ return x0
416
+
417
+ def model_fn(self, x, t):
418
+ """
419
+ Convert the model to the noise prediction model or the data prediction model.
420
+ """
421
+ if self.predict_x0:
422
+ return self.data_prediction_fn(x, t)
423
+ else:
424
+ return self.noise_prediction_fn(x, t)
425
+
426
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
427
+ """Compute the intermediate time steps for sampling.
428
+ """
429
+ if skip_type == 'logSNR':
430
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
431
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
432
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
433
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
434
+ elif skip_type == 'time_uniform':
435
+ return torch.linspace(t_T, t_0, N + 1).to(device)
436
+ elif skip_type == 'time_quadratic':
437
+ t_order = 2
438
+ t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
439
+ return t
440
+ else:
441
+ raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
442
+
443
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
444
+ """
445
+ Get the order of each step for sampling by the singlestep DPM-Solver.
446
+ """
447
+ if order == 3:
448
+ K = steps // 3 + 1
449
+ if steps % 3 == 0:
450
+ orders = [3,] * (K - 2) + [2, 1]
451
+ elif steps % 3 == 1:
452
+ orders = [3,] * (K - 1) + [1]
453
+ else:
454
+ orders = [3,] * (K - 1) + [2]
455
+ elif order == 2:
456
+ if steps % 2 == 0:
457
+ K = steps // 2
458
+ orders = [2,] * K
459
+ else:
460
+ K = steps // 2 + 1
461
+ orders = [2,] * (K - 1) + [1]
462
+ elif order == 1:
463
+ K = steps
464
+ orders = [1,] * steps
465
+ else:
466
+ raise ValueError("'order' must be '1' or '2' or '3'.")
467
+ if skip_type == 'logSNR':
468
+ # To reproduce the results in DPM-Solver paper
469
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
470
+ else:
471
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
472
+ return timesteps_outer, orders
473
+
474
+ def denoise_to_zero_fn(self, x, s):
475
+ """
476
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
477
+ """
478
+ return self.data_prediction_fn(x, s)
479
+
480
+ def multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
481
+ if len(t.shape) == 0:
482
+ t = t.view(-1)
483
+ if 'bh' in self.variant:
484
+ return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
485
+ else:
486
+ assert self.variant == 'vary_coeff'
487
+ return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
488
+
489
+ def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
490
+ print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
491
+ ns = self.noise_schedule
492
+ assert order <= len(model_prev_list)
493
+
494
+ # first compute rks
495
+ t_prev_0 = t_prev_list[-1]
496
+ lambda_prev_0 = ns.marginal_lambda(t_prev_0)
497
+ lambda_t = ns.marginal_lambda(t)
498
+ model_prev_0 = model_prev_list[-1]
499
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
500
+ log_alpha_t = ns.marginal_log_mean_coeff(t)
501
+ alpha_t = torch.exp(log_alpha_t)
502
+
503
+ h = lambda_t - lambda_prev_0
504
+
505
+ rks = []
506
+ D1s = []
507
+ for i in range(1, order):
508
+ t_prev_i = t_prev_list[-(i + 1)]
509
+ model_prev_i = model_prev_list[-(i + 1)]
510
+ lambda_prev_i = ns.marginal_lambda(t_prev_i)
511
+ rk = (lambda_prev_i - lambda_prev_0) / h
512
+ rks.append(rk)
513
+ D1s.append((model_prev_i - model_prev_0) / rk)
514
+
515
+ rks.append(1.)
516
+ rks = torch.tensor(rks, device=x.device)
517
+
518
+ K = len(rks)
519
+ # build C matrix
520
+ C = []
521
+
522
+ col = torch.ones_like(rks)
523
+ for k in range(1, K + 1):
524
+ C.append(col)
525
+ col = col * rks / (k + 1)
526
+ C = torch.stack(C, dim=1)
527
+
528
+ if len(D1s) > 0:
529
+ D1s = torch.stack(D1s, dim=1) # (B, K)
530
+ C_inv_p = torch.linalg.inv(C[:-1, :-1])
531
+ A_p = C_inv_p
532
+
533
+ if use_corrector:
534
+ print('using corrector')
535
+ C_inv = torch.linalg.inv(C)
536
+ A_c = C_inv
537
+
538
+ hh = -h if self.predict_x0 else h
539
+ h_phi_1 = torch.expm1(hh)
540
+ h_phi_ks = []
541
+ factorial_k = 1
542
+ h_phi_k = h_phi_1
543
+ for k in range(1, K + 2):
544
+ h_phi_ks.append(h_phi_k)
545
+ h_phi_k = h_phi_k / hh - 1 / factorial_k
546
+ factorial_k *= (k + 1)
547
+
548
+ model_t = None
549
+ if self.predict_x0:
550
+ x_t_ = (
551
+ sigma_t / sigma_prev_0 * x
552
+ - alpha_t * h_phi_1 * model_prev_0
553
+ )
554
+ # now predictor
555
+ x_t = x_t_
556
+ if len(D1s) > 0:
557
+ # compute the residuals for predictor
558
+ for k in range(K - 1):
559
+ x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
560
+ # now corrector
561
+ if use_corrector:
562
+ model_t = self.model_fn(x_t, t)
563
+ D1_t = (model_t - model_prev_0)
564
+ x_t = x_t_
565
+ k = 0
566
+ for k in range(K - 1):
567
+ x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
568
+ x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
569
+ else:
570
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
571
+ x_t_ = (
572
+ (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
573
+ - (sigma_t * h_phi_1) * model_prev_0
574
+ )
575
+ # now predictor
576
+ x_t = x_t_
577
+ if len(D1s) > 0:
578
+ # compute the residuals for predictor
579
+ for k in range(K - 1):
580
+ x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
581
+ # now corrector
582
+ if use_corrector:
583
+ model_t = self.model_fn(x_t, t)
584
+ D1_t = (model_t - model_prev_0)
585
+ x_t = x_t_
586
+ k = 0
587
+ for k in range(K - 1):
588
+ x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
589
+ x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
590
+ return x_t, model_t
591
+
592
+ def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
593
+ # print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
594
+ ns = self.noise_schedule
595
+ assert order <= len(model_prev_list)
596
+ dims = x.dim()
597
+
598
+ # first compute rks
599
+ t_prev_0 = t_prev_list[-1]
600
+ lambda_prev_0 = ns.marginal_lambda(t_prev_0)
601
+ lambda_t = ns.marginal_lambda(t)
602
+ model_prev_0 = model_prev_list[-1]
603
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
604
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
605
+ alpha_t = torch.exp(log_alpha_t)
606
+
607
+ h = lambda_t - lambda_prev_0
608
+
609
+ rks = []
610
+ D1s = []
611
+ for i in range(1, order):
612
+ t_prev_i = t_prev_list[-(i + 1)]
613
+ model_prev_i = model_prev_list[-(i + 1)]
614
+ lambda_prev_i = ns.marginal_lambda(t_prev_i)
615
+ rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
616
+ rks.append(rk)
617
+ D1s.append((model_prev_i - model_prev_0) / rk)
618
+
619
+ rks.append(1.)
620
+ rks = torch.tensor(rks, device=x.device)
621
+
622
+ R = []
623
+ b = []
624
+
625
+ hh = -h[0] if self.predict_x0 else h[0]
626
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
627
+ h_phi_k = h_phi_1 / hh - 1
628
+
629
+ factorial_i = 1
630
+
631
+ if self.variant == 'bh1':
632
+ B_h = hh
633
+ elif self.variant == 'bh2':
634
+ B_h = torch.expm1(hh)
635
+ else:
636
+ raise NotImplementedError()
637
+
638
+ for i in range(1, order + 1):
639
+ R.append(torch.pow(rks, i - 1))
640
+ b.append(h_phi_k * factorial_i / B_h)
641
+ factorial_i *= (i + 1)
642
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
643
+
644
+ R = torch.stack(R)
645
+ b = torch.tensor(b, device=x.device)
646
+
647
+ # now predictor
648
+ use_predictor = len(D1s) > 0 and x_t is None
649
+ if len(D1s) > 0:
650
+ D1s = torch.stack(D1s, dim=1) # (B, K)
651
+ if x_t is None:
652
+ # for order 2, we use a simplified version
653
+ if order == 2:
654
+ rhos_p = torch.tensor([0.5], device=b.device)
655
+ else:
656
+ rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
657
+ else:
658
+ D1s = None
659
+
660
+ if use_corrector:
661
+ # print('using corrector')
662
+ # for order 1, we use a simplified version
663
+ if order == 1:
664
+ rhos_c = torch.tensor([0.5], device=b.device)
665
+ else:
666
+ rhos_c = torch.linalg.solve(R, b)
667
+
668
+ model_t = None
669
+ if self.predict_x0:
670
+ x_t_ = (
671
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
672
+ - expand_dims(alpha_t * h_phi_1, dims)* model_prev_0
673
+ )
674
+
675
+ if x_t is None:
676
+ if use_predictor:
677
+ pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
678
+ else:
679
+ pred_res = 0
680
+ x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * pred_res
681
+
682
+ if use_corrector:
683
+ model_t = self.model_fn(x_t, t)
684
+ if D1s is not None:
685
+ corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
686
+ else:
687
+ corr_res = 0
688
+ D1_t = (model_t - model_prev_0)
689
+ x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
690
+ else:
691
+ x_t_ = (
692
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dimss) * x
693
+ - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
694
+ )
695
+ if x_t is None:
696
+ if use_predictor:
697
+ pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
698
+ else:
699
+ pred_res = 0
700
+ x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * pred_res
701
+
702
+ if use_corrector:
703
+ model_t = self.model_fn(x_t, t)
704
+ if D1s is not None:
705
+ corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
706
+ else:
707
+ corr_res = 0
708
+ D1_t = (model_t - model_prev_0)
709
+ x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
710
+ return x_t, model_t
711
+
712
+
713
+ def sample(self, x, timesteps, t_start=None, t_end=None, order=3, skip_type='time_uniform',
714
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
715
+ atol=0.0078, rtol=0.05, corrector=False,
716
+ ):
717
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
718
+ t_T = self.noise_schedule.T if t_start is None else t_start
719
+ device = x.device
720
+ steps = len(timesteps) - 1
721
+ if method == 'multistep':
722
+ assert steps >= order
723
+ # timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
724
+ assert timesteps.shape[0] - 1 == steps
725
+ # with torch.no_grad():
726
+ for step_index in trange(steps):
727
+ if self.noise_mask is not None:
728
+ x = x * self.noise_mask + (1. - self.noise_mask) * (self.masked_image * self.noise_schedule.marginal_alpha(timesteps[step_index]) + self.noise * self.noise_schedule.marginal_std(timesteps[step_index]))
729
+ if step_index == 0:
730
+ vec_t = timesteps[0].expand((x.shape[0]))
731
+ model_prev_list = [self.model_fn(x, vec_t)]
732
+ t_prev_list = [vec_t]
733
+ elif step_index < order:
734
+ init_order = step_index
735
+ # Init the first `order` values by lower order multistep DPM-Solver.
736
+ # for init_order in range(1, order):
737
+ vec_t = timesteps[init_order].expand(x.shape[0])
738
+ x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
739
+ if model_x is None:
740
+ model_x = self.model_fn(x, vec_t)
741
+ model_prev_list.append(model_x)
742
+ t_prev_list.append(vec_t)
743
+ else:
744
+ extra_final_step = 0
745
+ if step_index == (steps - 1):
746
+ extra_final_step = 1
747
+ for step in range(step_index, step_index + 1 + extra_final_step):
748
+ vec_t = timesteps[step].expand(x.shape[0])
749
+ if lower_order_final:
750
+ step_order = min(order, steps + 1 - step)
751
+ else:
752
+ step_order = order
753
+ # print('this step order:', step_order)
754
+ if step == steps:
755
+ # print('do not run corrector at the last step')
756
+ use_corrector = False
757
+ else:
758
+ use_corrector = True
759
+ x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
760
+ for i in range(order - 1):
761
+ t_prev_list[i] = t_prev_list[i + 1]
762
+ model_prev_list[i] = model_prev_list[i + 1]
763
+ t_prev_list[-1] = vec_t
764
+ # We do not need to evaluate the final model value.
765
+ if step < steps:
766
+ if model_x is None:
767
+ model_x = self.model_fn(x, vec_t)
768
+ model_prev_list[-1] = model_x
769
+ else:
770
+ raise NotImplementedError()
771
+ if denoise_to_zero:
772
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
773
+ return x
774
+
775
+
776
+ #############################################################
777
+ # other utility functions
778
+ #############################################################
779
+
780
+ def interpolate_fn(x, xp, yp):
781
+ """
782
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
783
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
784
+ 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.)
785
+
786
+ Args:
787
+ 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).
788
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
789
+ yp: PyTorch tensor with shape [C, K].
790
+ Returns:
791
+ The function values f(x), with shape [N, C].
792
+ """
793
+ N, K = x.shape[0], xp.shape[1]
794
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
795
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
796
+ x_idx = torch.argmin(x_indices, dim=2)
797
+ cand_start_idx = x_idx - 1
798
+ start_idx = torch.where(
799
+ torch.eq(x_idx, 0),
800
+ torch.tensor(1, device=x.device),
801
+ torch.where(
802
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
803
+ ),
804
+ )
805
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
806
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
807
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
808
+ start_idx2 = torch.where(
809
+ torch.eq(x_idx, 0),
810
+ torch.tensor(0, device=x.device),
811
+ torch.where(
812
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
813
+ ),
814
+ )
815
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
816
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
817
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
818
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
819
+ return cand
820
+
821
+
822
+ def expand_dims(v, dims):
823
+ """
824
+ Expand the tensor `v` to the dim `dims`.
825
+
826
+ Args:
827
+ `v`: a PyTorch tensor with shape [N].
828
+ `dim`: a `int`.
829
+ Returns:
830
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
831
+ """
832
+ return v[(...,) + (None,)*(dims - 1)]
833
+
834
+
835
+
836
+ def sample_unipc(model, noise, image, sigmas, sampling_function, max_denoise, extra_args=None, callback=None, disable=None, noise_mask=None, variant='bh1'):
837
+ to_zero = False
838
+ if sigmas[-1] == 0:
839
+ timesteps = torch.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
840
+ to_zero = True
841
+ else:
842
+ timesteps = sigmas.clone()
843
+
844
+ for s in range(timesteps.shape[0]):
845
+ timesteps[s] = (model.sigma_to_t(timesteps[s]) / 1000) + (1 / len(model.sigmas))
846
+
847
+ ns = NoiseScheduleVP('discrete', alphas_cumprod=model.inner_model.alphas_cumprod)
848
+
849
+ if image is not None:
850
+ img = image * ns.marginal_alpha(timesteps[0])
851
+ if max_denoise:
852
+ noise_mult = 1.0
853
+ else:
854
+ noise_mult = ns.marginal_std(timesteps[0])
855
+ img += noise * noise_mult
856
+ else:
857
+ img = noise
858
+
859
+ if to_zero:
860
+ timesteps[-1] = (1 / len(model.sigmas))
861
+
862
+ device = noise.device
863
+
864
+ if model.parameterization == "v":
865
+ model_type = "v"
866
+ else:
867
+ model_type = "noise"
868
+
869
+ model_fn = model_wrapper(
870
+ model.inner_model.inner_model.apply_model,
871
+ sampling_function,
872
+ ns,
873
+ model_type=model_type,
874
+ guidance_type="uncond",
875
+ model_kwargs=extra_args,
876
+ )
877
+
878
+ order = min(3, len(timesteps) - 1)
879
+ uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, noise_mask=noise_mask, masked_image=image, noise=noise, variant=variant)
880
+ x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True)
881
+ if not to_zero:
882
+ x /= ns.marginal_alpha(timesteps[-1])
883
+ return x
comfy/k_diffusion/augmentation.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from functools import reduce
2
+ import math
3
+ import operator
4
+
5
+ import numpy as np
6
+ from skimage import transform
7
+ import torch
8
+ from torch import nn
9
+
10
+
11
+ def translate2d(tx, ty):
12
+ mat = [[1, 0, tx],
13
+ [0, 1, ty],
14
+ [0, 0, 1]]
15
+ return torch.tensor(mat, dtype=torch.float32)
16
+
17
+
18
+ def scale2d(sx, sy):
19
+ mat = [[sx, 0, 0],
20
+ [ 0, sy, 0],
21
+ [ 0, 0, 1]]
22
+ return torch.tensor(mat, dtype=torch.float32)
23
+
24
+
25
+ def rotate2d(theta):
26
+ mat = [[torch.cos(theta), torch.sin(-theta), 0],
27
+ [torch.sin(theta), torch.cos(theta), 0],
28
+ [ 0, 0, 1]]
29
+ return torch.tensor(mat, dtype=torch.float32)
30
+
31
+
32
+ class KarrasAugmentationPipeline:
33
+ def __init__(self, a_prob=0.12, a_scale=2**0.2, a_aniso=2**0.2, a_trans=1/8):
34
+ self.a_prob = a_prob
35
+ self.a_scale = a_scale
36
+ self.a_aniso = a_aniso
37
+ self.a_trans = a_trans
38
+
39
+ def __call__(self, image):
40
+ h, w = image.size
41
+ mats = [translate2d(h / 2 - 0.5, w / 2 - 0.5)]
42
+
43
+ # x-flip
44
+ a0 = torch.randint(2, []).float()
45
+ mats.append(scale2d(1 - 2 * a0, 1))
46
+ # y-flip
47
+ do = (torch.rand([]) < self.a_prob).float()
48
+ a1 = torch.randint(2, []).float() * do
49
+ mats.append(scale2d(1, 1 - 2 * a1))
50
+ # scaling
51
+ do = (torch.rand([]) < self.a_prob).float()
52
+ a2 = torch.randn([]) * do
53
+ mats.append(scale2d(self.a_scale ** a2, self.a_scale ** a2))
54
+ # rotation
55
+ do = (torch.rand([]) < self.a_prob).float()
56
+ a3 = (torch.rand([]) * 2 * math.pi - math.pi) * do
57
+ mats.append(rotate2d(-a3))
58
+ # anisotropy
59
+ do = (torch.rand([]) < self.a_prob).float()
60
+ a4 = (torch.rand([]) * 2 * math.pi - math.pi) * do
61
+ a5 = torch.randn([]) * do
62
+ mats.append(rotate2d(a4))
63
+ mats.append(scale2d(self.a_aniso ** a5, self.a_aniso ** -a5))
64
+ mats.append(rotate2d(-a4))
65
+ # translation
66
+ do = (torch.rand([]) < self.a_prob).float()
67
+ a6 = torch.randn([]) * do
68
+ a7 = torch.randn([]) * do
69
+ mats.append(translate2d(self.a_trans * w * a6, self.a_trans * h * a7))
70
+
71
+ # form the transformation matrix and conditioning vector
72
+ mats.append(translate2d(-h / 2 + 0.5, -w / 2 + 0.5))
73
+ mat = reduce(operator.matmul, mats)
74
+ cond = torch.stack([a0, a1, a2, a3.cos() - 1, a3.sin(), a5 * a4.cos(), a5 * a4.sin(), a6, a7])
75
+
76
+ # apply the transformation
77
+ image_orig = np.array(image, dtype=np.float32) / 255
78
+ if image_orig.ndim == 2:
79
+ image_orig = image_orig[..., None]
80
+ tf = transform.AffineTransform(mat.numpy())
81
+ image = transform.warp(image_orig, tf.inverse, order=3, mode='reflect', cval=0.5, clip=False, preserve_range=True)
82
+ image_orig = torch.as_tensor(image_orig).movedim(2, 0) * 2 - 1
83
+ image = torch.as_tensor(image).movedim(2, 0) * 2 - 1
84
+ return image, image_orig, cond
85
+
86
+
87
+ class KarrasAugmentWrapper(nn.Module):
88
+ def __init__(self, model):
89
+ super().__init__()
90
+ self.inner_model = model
91
+
92
+ def forward(self, input, sigma, aug_cond=None, mapping_cond=None, **kwargs):
93
+ if aug_cond is None:
94
+ aug_cond = input.new_zeros([input.shape[0], 9])
95
+ if mapping_cond is None:
96
+ mapping_cond = aug_cond
97
+ else:
98
+ mapping_cond = torch.cat([aug_cond, mapping_cond], dim=1)
99
+ return self.inner_model(input, sigma, mapping_cond=mapping_cond, **kwargs)
100
+
101
+ def set_skip_stages(self, skip_stages):
102
+ return self.inner_model.set_skip_stages(skip_stages)
103
+
104
+ def set_patch_size(self, patch_size):
105
+ return self.inner_model.set_patch_size(patch_size)
comfy/k_diffusion/config.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from functools import partial
2
+ import json
3
+ import math
4
+ import warnings
5
+
6
+ from jsonmerge import merge
7
+
8
+ from . import augmentation, layers, models, utils
9
+
10
+
11
+ def load_config(file):
12
+ defaults = {
13
+ 'model': {
14
+ 'sigma_data': 1.,
15
+ 'patch_size': 1,
16
+ 'dropout_rate': 0.,
17
+ 'augment_wrapper': True,
18
+ 'augment_prob': 0.,
19
+ 'mapping_cond_dim': 0,
20
+ 'unet_cond_dim': 0,
21
+ 'cross_cond_dim': 0,
22
+ 'cross_attn_depths': None,
23
+ 'skip_stages': 0,
24
+ 'has_variance': False,
25
+ },
26
+ 'dataset': {
27
+ 'type': 'imagefolder',
28
+ },
29
+ 'optimizer': {
30
+ 'type': 'adamw',
31
+ 'lr': 1e-4,
32
+ 'betas': [0.95, 0.999],
33
+ 'eps': 1e-6,
34
+ 'weight_decay': 1e-3,
35
+ },
36
+ 'lr_sched': {
37
+ 'type': 'inverse',
38
+ 'inv_gamma': 20000.,
39
+ 'power': 1.,
40
+ 'warmup': 0.99,
41
+ },
42
+ 'ema_sched': {
43
+ 'type': 'inverse',
44
+ 'power': 0.6667,
45
+ 'max_value': 0.9999
46
+ },
47
+ }
48
+ config = json.load(file)
49
+ return merge(defaults, config)
50
+
51
+
52
+ def make_model(config):
53
+ config = config['model']
54
+ assert config['type'] == 'image_v1'
55
+ model = models.ImageDenoiserModelV1(
56
+ config['input_channels'],
57
+ config['mapping_out'],
58
+ config['depths'],
59
+ config['channels'],
60
+ config['self_attn_depths'],
61
+ config['cross_attn_depths'],
62
+ patch_size=config['patch_size'],
63
+ dropout_rate=config['dropout_rate'],
64
+ mapping_cond_dim=config['mapping_cond_dim'] + (9 if config['augment_wrapper'] else 0),
65
+ unet_cond_dim=config['unet_cond_dim'],
66
+ cross_cond_dim=config['cross_cond_dim'],
67
+ skip_stages=config['skip_stages'],
68
+ has_variance=config['has_variance'],
69
+ )
70
+ if config['augment_wrapper']:
71
+ model = augmentation.KarrasAugmentWrapper(model)
72
+ return model
73
+
74
+
75
+ def make_denoiser_wrapper(config):
76
+ config = config['model']
77
+ sigma_data = config.get('sigma_data', 1.)
78
+ has_variance = config.get('has_variance', False)
79
+ if not has_variance:
80
+ return partial(layers.Denoiser, sigma_data=sigma_data)
81
+ return partial(layers.DenoiserWithVariance, sigma_data=sigma_data)
82
+
83
+
84
+ def make_sample_density(config):
85
+ sd_config = config['sigma_sample_density']
86
+ sigma_data = config['sigma_data']
87
+ if sd_config['type'] == 'lognormal':
88
+ loc = sd_config['mean'] if 'mean' in sd_config else sd_config['loc']
89
+ scale = sd_config['std'] if 'std' in sd_config else sd_config['scale']
90
+ return partial(utils.rand_log_normal, loc=loc, scale=scale)
91
+ if sd_config['type'] == 'loglogistic':
92
+ loc = sd_config['loc'] if 'loc' in sd_config else math.log(sigma_data)
93
+ scale = sd_config['scale'] if 'scale' in sd_config else 0.5
94
+ min_value = sd_config['min_value'] if 'min_value' in sd_config else 0.
95
+ max_value = sd_config['max_value'] if 'max_value' in sd_config else float('inf')
96
+ return partial(utils.rand_log_logistic, loc=loc, scale=scale, min_value=min_value, max_value=max_value)
97
+ if sd_config['type'] == 'loguniform':
98
+ min_value = sd_config['min_value'] if 'min_value' in sd_config else config['sigma_min']
99
+ max_value = sd_config['max_value'] if 'max_value' in sd_config else config['sigma_max']
100
+ return partial(utils.rand_log_uniform, min_value=min_value, max_value=max_value)
101
+ if sd_config['type'] == 'v-diffusion':
102
+ min_value = sd_config['min_value'] if 'min_value' in sd_config else 0.
103
+ max_value = sd_config['max_value'] if 'max_value' in sd_config else float('inf')
104
+ return partial(utils.rand_v_diffusion, sigma_data=sigma_data, min_value=min_value, max_value=max_value)
105
+ if sd_config['type'] == 'split-lognormal':
106
+ loc = sd_config['mean'] if 'mean' in sd_config else sd_config['loc']
107
+ scale_1 = sd_config['std_1'] if 'std_1' in sd_config else sd_config['scale_1']
108
+ scale_2 = sd_config['std_2'] if 'std_2' in sd_config else sd_config['scale_2']
109
+ return partial(utils.rand_split_log_normal, loc=loc, scale_1=scale_1, scale_2=scale_2)
110
+ raise ValueError('Unknown sample density type')
comfy/k_diffusion/evaluation.py ADDED
@@ -0,0 +1,134 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import os
3
+ from pathlib import Path
4
+
5
+ from cleanfid.inception_torchscript import InceptionV3W
6
+ import clip
7
+ from resize_right import resize
8
+ import torch
9
+ from torch import nn
10
+ from torch.nn import functional as F
11
+ from torchvision import transforms
12
+ from tqdm.auto import trange
13
+
14
+ from . import utils
15
+
16
+
17
+ class InceptionV3FeatureExtractor(nn.Module):
18
+ def __init__(self, device='cpu'):
19
+ super().__init__()
20
+ path = Path(os.environ.get('XDG_CACHE_HOME', Path.home() / '.cache')) / 'k-diffusion'
21
+ url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
22
+ digest = 'f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4'
23
+ utils.download_file(path / 'inception-2015-12-05.pt', url, digest)
24
+ self.model = InceptionV3W(str(path), resize_inside=False).to(device)
25
+ self.size = (299, 299)
26
+
27
+ def forward(self, x):
28
+ if x.shape[2:4] != self.size:
29
+ x = resize(x, out_shape=self.size, pad_mode='reflect')
30
+ if x.shape[1] == 1:
31
+ x = torch.cat([x] * 3, dim=1)
32
+ x = (x * 127.5 + 127.5).clamp(0, 255)
33
+ return self.model(x)
34
+
35
+
36
+ class CLIPFeatureExtractor(nn.Module):
37
+ def __init__(self, name='ViT-L/14@336px', device='cpu'):
38
+ super().__init__()
39
+ self.model = clip.load(name, device=device)[0].eval().requires_grad_(False)
40
+ self.normalize = transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073),
41
+ std=(0.26862954, 0.26130258, 0.27577711))
42
+ self.size = (self.model.visual.input_resolution, self.model.visual.input_resolution)
43
+
44
+ def forward(self, x):
45
+ if x.shape[2:4] != self.size:
46
+ x = resize(x.add(1).div(2), out_shape=self.size, pad_mode='reflect').clamp(0, 1)
47
+ x = self.normalize(x)
48
+ x = self.model.encode_image(x).float()
49
+ x = F.normalize(x) * x.shape[1] ** 0.5
50
+ return x
51
+
52
+
53
+ def compute_features(accelerator, sample_fn, extractor_fn, n, batch_size):
54
+ n_per_proc = math.ceil(n / accelerator.num_processes)
55
+ feats_all = []
56
+ try:
57
+ for i in trange(0, n_per_proc, batch_size, disable=not accelerator.is_main_process):
58
+ cur_batch_size = min(n - i, batch_size)
59
+ samples = sample_fn(cur_batch_size)[:cur_batch_size]
60
+ feats_all.append(accelerator.gather(extractor_fn(samples)))
61
+ except StopIteration:
62
+ pass
63
+ return torch.cat(feats_all)[:n]
64
+
65
+
66
+ def polynomial_kernel(x, y):
67
+ d = x.shape[-1]
68
+ dot = x @ y.transpose(-2, -1)
69
+ return (dot / d + 1) ** 3
70
+
71
+
72
+ def squared_mmd(x, y, kernel=polynomial_kernel):
73
+ m = x.shape[-2]
74
+ n = y.shape[-2]
75
+ kxx = kernel(x, x)
76
+ kyy = kernel(y, y)
77
+ kxy = kernel(x, y)
78
+ kxx_sum = kxx.sum([-1, -2]) - kxx.diagonal(dim1=-1, dim2=-2).sum(-1)
79
+ kyy_sum = kyy.sum([-1, -2]) - kyy.diagonal(dim1=-1, dim2=-2).sum(-1)
80
+ kxy_sum = kxy.sum([-1, -2])
81
+ term_1 = kxx_sum / m / (m - 1)
82
+ term_2 = kyy_sum / n / (n - 1)
83
+ term_3 = kxy_sum * 2 / m / n
84
+ return term_1 + term_2 - term_3
85
+
86
+
87
+ @utils.tf32_mode(matmul=False)
88
+ def kid(x, y, max_size=5000):
89
+ x_size, y_size = x.shape[0], y.shape[0]
90
+ n_partitions = math.ceil(max(x_size / max_size, y_size / max_size))
91
+ total_mmd = x.new_zeros([])
92
+ for i in range(n_partitions):
93
+ cur_x = x[round(i * x_size / n_partitions):round((i + 1) * x_size / n_partitions)]
94
+ cur_y = y[round(i * y_size / n_partitions):round((i + 1) * y_size / n_partitions)]
95
+ total_mmd = total_mmd + squared_mmd(cur_x, cur_y)
96
+ return total_mmd / n_partitions
97
+
98
+
99
+ class _MatrixSquareRootEig(torch.autograd.Function):
100
+ @staticmethod
101
+ def forward(ctx, a):
102
+ vals, vecs = torch.linalg.eigh(a)
103
+ ctx.save_for_backward(vals, vecs)
104
+ return vecs @ vals.abs().sqrt().diag_embed() @ vecs.transpose(-2, -1)
105
+
106
+ @staticmethod
107
+ def backward(ctx, grad_output):
108
+ vals, vecs = ctx.saved_tensors
109
+ d = vals.abs().sqrt().unsqueeze(-1).repeat_interleave(vals.shape[-1], -1)
110
+ vecs_t = vecs.transpose(-2, -1)
111
+ return vecs @ (vecs_t @ grad_output @ vecs / (d + d.transpose(-2, -1))) @ vecs_t
112
+
113
+
114
+ def sqrtm_eig(a):
115
+ if a.ndim < 2:
116
+ raise RuntimeError('tensor of matrices must have at least 2 dimensions')
117
+ if a.shape[-2] != a.shape[-1]:
118
+ raise RuntimeError('tensor must be batches of square matrices')
119
+ return _MatrixSquareRootEig.apply(a)
120
+
121
+
122
+ @utils.tf32_mode(matmul=False)
123
+ def fid(x, y, eps=1e-8):
124
+ x_mean = x.mean(dim=0)
125
+ y_mean = y.mean(dim=0)
126
+ mean_term = (x_mean - y_mean).pow(2).sum()
127
+ x_cov = torch.cov(x.T)
128
+ y_cov = torch.cov(y.T)
129
+ eps_eye = torch.eye(x_cov.shape[0], device=x_cov.device, dtype=x_cov.dtype) * eps
130
+ x_cov = x_cov + eps_eye
131
+ y_cov = y_cov + eps_eye
132
+ x_cov_sqrt = sqrtm_eig(x_cov)
133
+ cov_term = torch.trace(x_cov + y_cov - 2 * sqrtm_eig(x_cov_sqrt @ y_cov @ x_cov_sqrt))
134
+ return mean_term + cov_term
comfy/k_diffusion/external.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ from torch import nn
5
+
6
+ from . import sampling, utils
7
+
8
+
9
+ class VDenoiser(nn.Module):
10
+ """A v-diffusion-pytorch model wrapper for k-diffusion."""
11
+
12
+ def __init__(self, inner_model):
13
+ super().__init__()
14
+ self.inner_model = inner_model
15
+ self.sigma_data = 1.
16
+
17
+ def get_scalings(self, sigma):
18
+ c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
19
+ c_out = -sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
20
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
21
+ return c_skip, c_out, c_in
22
+
23
+ def sigma_to_t(self, sigma):
24
+ return sigma.atan() / math.pi * 2
25
+
26
+ def t_to_sigma(self, t):
27
+ return (t * math.pi / 2).tan()
28
+
29
+ def loss(self, input, noise, sigma, **kwargs):
30
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
31
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
32
+ model_output = self.inner_model(noised_input * c_in, self.sigma_to_t(sigma), **kwargs)
33
+ target = (input - c_skip * noised_input) / c_out
34
+ return (model_output - target).pow(2).flatten(1).mean(1)
35
+
36
+ def forward(self, input, sigma, **kwargs):
37
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
38
+ return self.inner_model(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip
39
+
40
+
41
+ class DiscreteSchedule(nn.Module):
42
+ """A mapping between continuous noise levels (sigmas) and a list of discrete noise
43
+ levels."""
44
+
45
+ def __init__(self, sigmas, quantize):
46
+ super().__init__()
47
+ self.register_buffer('sigmas', sigmas)
48
+ self.register_buffer('log_sigmas', sigmas.log())
49
+ self.quantize = quantize
50
+
51
+ @property
52
+ def sigma_min(self):
53
+ return self.sigmas[0]
54
+
55
+ @property
56
+ def sigma_max(self):
57
+ return self.sigmas[-1]
58
+
59
+ def get_sigmas(self, n=None):
60
+ if n is None:
61
+ return sampling.append_zero(self.sigmas.flip(0))
62
+ t_max = len(self.sigmas) - 1
63
+ t = torch.linspace(t_max, 0, n, device=self.sigmas.device)
64
+ return sampling.append_zero(self.t_to_sigma(t))
65
+
66
+ def sigma_to_t(self, sigma, quantize=None):
67
+ quantize = self.quantize if quantize is None else quantize
68
+ log_sigma = sigma.log()
69
+ dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
70
+ if quantize:
71
+ return dists.abs().argmin(dim=0).view(sigma.shape)
72
+ low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2)
73
+ high_idx = low_idx + 1
74
+ low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx]
75
+ w = (low - log_sigma) / (low - high)
76
+ w = w.clamp(0, 1)
77
+ t = (1 - w) * low_idx + w * high_idx
78
+ return t.view(sigma.shape)
79
+
80
+ def t_to_sigma(self, t):
81
+ t = t.float()
82
+ low_idx = t.floor().long()
83
+ high_idx = t.ceil().long()
84
+ w = t-low_idx if t.device.type == 'mps' else t.frac()
85
+ log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
86
+ return log_sigma.exp()
87
+
88
+
89
+ class DiscreteEpsDDPMDenoiser(DiscreteSchedule):
90
+ """A wrapper for discrete schedule DDPM models that output eps (the predicted
91
+ noise)."""
92
+
93
+ def __init__(self, model, alphas_cumprod, quantize):
94
+ super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize)
95
+ self.inner_model = model
96
+ self.sigma_data = 1.
97
+
98
+ def get_scalings(self, sigma):
99
+ c_out = -sigma
100
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
101
+ return c_out, c_in
102
+
103
+ def get_eps(self, *args, **kwargs):
104
+ return self.inner_model(*args, **kwargs)
105
+
106
+ def loss(self, input, noise, sigma, **kwargs):
107
+ c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
108
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
109
+ eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs)
110
+ return (eps - noise).pow(2).flatten(1).mean(1)
111
+
112
+ def forward(self, input, sigma, **kwargs):
113
+ c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
114
+ eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
115
+ return input + eps * c_out
116
+
117
+
118
+ class OpenAIDenoiser(DiscreteEpsDDPMDenoiser):
119
+ """A wrapper for OpenAI diffusion models."""
120
+
121
+ def __init__(self, model, diffusion, quantize=False, has_learned_sigmas=True, device='cpu'):
122
+ alphas_cumprod = torch.tensor(diffusion.alphas_cumprod, device=device, dtype=torch.float32)
123
+ super().__init__(model, alphas_cumprod, quantize=quantize)
124
+ self.has_learned_sigmas = has_learned_sigmas
125
+
126
+ def get_eps(self, *args, **kwargs):
127
+ model_output = self.inner_model(*args, **kwargs)
128
+ if self.has_learned_sigmas:
129
+ return model_output.chunk(2, dim=1)[0]
130
+ return model_output
131
+
132
+
133
+ class CompVisDenoiser(DiscreteEpsDDPMDenoiser):
134
+ """A wrapper for CompVis diffusion models."""
135
+
136
+ def __init__(self, model, quantize=False, device='cpu'):
137
+ super().__init__(model, model.alphas_cumprod, quantize=quantize)
138
+
139
+ def get_eps(self, *args, **kwargs):
140
+ return self.inner_model.apply_model(*args, **kwargs)
141
+
142
+
143
+ class DiscreteVDDPMDenoiser(DiscreteSchedule):
144
+ """A wrapper for discrete schedule DDPM models that output v."""
145
+
146
+ def __init__(self, model, alphas_cumprod, quantize):
147
+ super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize)
148
+ self.inner_model = model
149
+ self.sigma_data = 1.
150
+
151
+ def get_scalings(self, sigma):
152
+ c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
153
+ c_out = -sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
154
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
155
+ return c_skip, c_out, c_in
156
+
157
+ def get_v(self, *args, **kwargs):
158
+ return self.inner_model(*args, **kwargs)
159
+
160
+ def loss(self, input, noise, sigma, **kwargs):
161
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
162
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
163
+ model_output = self.get_v(noised_input * c_in, self.sigma_to_t(sigma), **kwargs)
164
+ target = (input - c_skip * noised_input) / c_out
165
+ return (model_output - target).pow(2).flatten(1).mean(1)
166
+
167
+ def forward(self, input, sigma, **kwargs):
168
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
169
+ return self.get_v(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip
170
+
171
+
172
+ class CompVisVDenoiser(DiscreteVDDPMDenoiser):
173
+ """A wrapper for CompVis diffusion models that output v."""
174
+
175
+ def __init__(self, model, quantize=False, device='cpu'):
176
+ super().__init__(model, model.alphas_cumprod, quantize=quantize)
177
+
178
+ def get_v(self, x, t, cond, **kwargs):
179
+ return self.inner_model.apply_model(x, t, cond)
comfy/k_diffusion/gns.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn
3
+
4
+
5
+ class DDPGradientStatsHook:
6
+ def __init__(self, ddp_module):
7
+ try:
8
+ ddp_module.register_comm_hook(self, self._hook_fn)
9
+ except AttributeError:
10
+ raise ValueError('DDPGradientStatsHook does not support non-DDP wrapped modules')
11
+ self._clear_state()
12
+
13
+ def _clear_state(self):
14
+ self.bucket_sq_norms_small_batch = []
15
+ self.bucket_sq_norms_large_batch = []
16
+
17
+ @staticmethod
18
+ def _hook_fn(self, bucket):
19
+ buf = bucket.buffer()
20
+ self.bucket_sq_norms_small_batch.append(buf.pow(2).sum())
21
+ fut = torch.distributed.all_reduce(buf, op=torch.distributed.ReduceOp.AVG, async_op=True).get_future()
22
+ def callback(fut):
23
+ buf = fut.value()[0]
24
+ self.bucket_sq_norms_large_batch.append(buf.pow(2).sum())
25
+ return buf
26
+ return fut.then(callback)
27
+
28
+ def get_stats(self):
29
+ sq_norm_small_batch = sum(self.bucket_sq_norms_small_batch)
30
+ sq_norm_large_batch = sum(self.bucket_sq_norms_large_batch)
31
+ self._clear_state()
32
+ stats = torch.stack([sq_norm_small_batch, sq_norm_large_batch])
33
+ torch.distributed.all_reduce(stats, op=torch.distributed.ReduceOp.AVG)
34
+ return stats[0].item(), stats[1].item()
35
+
36
+
37
+ class GradientNoiseScale:
38
+ """Calculates the gradient noise scale (1 / SNR), or critical batch size,
39
+ from _An Empirical Model of Large-Batch Training_,
40
+ https://arxiv.org/abs/1812.06162).
41
+
42
+ Args:
43
+ beta (float): The decay factor for the exponential moving averages used to
44
+ calculate the gradient noise scale.
45
+ Default: 0.9998
46
+ eps (float): Added for numerical stability.
47
+ Default: 1e-8
48
+ """
49
+
50
+ def __init__(self, beta=0.9998, eps=1e-8):
51
+ self.beta = beta
52
+ self.eps = eps
53
+ self.ema_sq_norm = 0.
54
+ self.ema_var = 0.
55
+ self.beta_cumprod = 1.
56
+ self.gradient_noise_scale = float('nan')
57
+
58
+ def state_dict(self):
59
+ """Returns the state of the object as a :class:`dict`."""
60
+ return dict(self.__dict__.items())
61
+
62
+ def load_state_dict(self, state_dict):
63
+ """Loads the object's state.
64
+ Args:
65
+ state_dict (dict): object state. Should be an object returned
66
+ from a call to :meth:`state_dict`.
67
+ """
68
+ self.__dict__.update(state_dict)
69
+
70
+ def update(self, sq_norm_small_batch, sq_norm_large_batch, n_small_batch, n_large_batch):
71
+ """Updates the state with a new batch's gradient statistics, and returns the
72
+ current gradient noise scale.
73
+
74
+ Args:
75
+ sq_norm_small_batch (float): The mean of the squared 2-norms of microbatch or
76
+ per sample gradients.
77
+ sq_norm_large_batch (float): The squared 2-norm of the mean of the microbatch or
78
+ per sample gradients.
79
+ n_small_batch (int): The batch size of the individual microbatch or per sample
80
+ gradients (1 if per sample).
81
+ n_large_batch (int): The total batch size of the mean of the microbatch or
82
+ per sample gradients.
83
+ """
84
+ est_sq_norm = (n_large_batch * sq_norm_large_batch - n_small_batch * sq_norm_small_batch) / (n_large_batch - n_small_batch)
85
+ est_var = (sq_norm_small_batch - sq_norm_large_batch) / (1 / n_small_batch - 1 / n_large_batch)
86
+ self.ema_sq_norm = self.beta * self.ema_sq_norm + (1 - self.beta) * est_sq_norm
87
+ self.ema_var = self.beta * self.ema_var + (1 - self.beta) * est_var
88
+ self.beta_cumprod *= self.beta
89
+ self.gradient_noise_scale = max(self.ema_var, self.eps) / max(self.ema_sq_norm, self.eps)
90
+ return self.gradient_noise_scale
91
+
92
+ def get_gns(self):
93
+ """Returns the current gradient noise scale."""
94
+ return self.gradient_noise_scale
95
+
96
+ def get_stats(self):
97
+ """Returns the current (debiased) estimates of the squared mean gradient
98
+ and gradient variance."""
99
+ return self.ema_sq_norm / (1 - self.beta_cumprod), self.ema_var / (1 - self.beta_cumprod)
comfy/k_diffusion/layers.py ADDED
@@ -0,0 +1,246 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ from einops import rearrange, repeat
4
+ import torch
5
+ from torch import nn
6
+ from torch.nn import functional as F
7
+
8
+ from . import utils
9
+
10
+ # Karras et al. preconditioned denoiser
11
+
12
+ class Denoiser(nn.Module):
13
+ """A Karras et al. preconditioner for denoising diffusion models."""
14
+
15
+ def __init__(self, inner_model, sigma_data=1.):
16
+ super().__init__()
17
+ self.inner_model = inner_model
18
+ self.sigma_data = sigma_data
19
+
20
+ def get_scalings(self, sigma):
21
+ c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
22
+ c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
23
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
24
+ return c_skip, c_out, c_in
25
+
26
+ def loss(self, input, noise, sigma, **kwargs):
27
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
28
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
29
+ model_output = self.inner_model(noised_input * c_in, sigma, **kwargs)
30
+ target = (input - c_skip * noised_input) / c_out
31
+ return (model_output - target).pow(2).flatten(1).mean(1)
32
+
33
+ def forward(self, input, sigma, **kwargs):
34
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
35
+ return self.inner_model(input * c_in, sigma, **kwargs) * c_out + input * c_skip
36
+
37
+
38
+ class DenoiserWithVariance(Denoiser):
39
+ def loss(self, input, noise, sigma, **kwargs):
40
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
41
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
42
+ model_output, logvar = self.inner_model(noised_input * c_in, sigma, return_variance=True, **kwargs)
43
+ logvar = utils.append_dims(logvar, model_output.ndim)
44
+ target = (input - c_skip * noised_input) / c_out
45
+ losses = ((model_output - target) ** 2 / logvar.exp() + logvar) / 2
46
+ return losses.flatten(1).mean(1)
47
+
48
+
49
+ # Residual blocks
50
+
51
+ class ResidualBlock(nn.Module):
52
+ def __init__(self, *main, skip=None):
53
+ super().__init__()
54
+ self.main = nn.Sequential(*main)
55
+ self.skip = skip if skip else nn.Identity()
56
+
57
+ def forward(self, input):
58
+ return self.main(input) + self.skip(input)
59
+
60
+
61
+ # Noise level (and other) conditioning
62
+
63
+ class ConditionedModule(nn.Module):
64
+ pass
65
+
66
+
67
+ class UnconditionedModule(ConditionedModule):
68
+ def __init__(self, module):
69
+ super().__init__()
70
+ self.module = module
71
+
72
+ def forward(self, input, cond=None):
73
+ return self.module(input)
74
+
75
+
76
+ class ConditionedSequential(nn.Sequential, ConditionedModule):
77
+ def forward(self, input, cond):
78
+ for module in self:
79
+ if isinstance(module, ConditionedModule):
80
+ input = module(input, cond)
81
+ else:
82
+ input = module(input)
83
+ return input
84
+
85
+
86
+ class ConditionedResidualBlock(ConditionedModule):
87
+ def __init__(self, *main, skip=None):
88
+ super().__init__()
89
+ self.main = ConditionedSequential(*main)
90
+ self.skip = skip if skip else nn.Identity()
91
+
92
+ def forward(self, input, cond):
93
+ skip = self.skip(input, cond) if isinstance(self.skip, ConditionedModule) else self.skip(input)
94
+ return self.main(input, cond) + skip
95
+
96
+
97
+ class AdaGN(ConditionedModule):
98
+ def __init__(self, feats_in, c_out, num_groups, eps=1e-5, cond_key='cond'):
99
+ super().__init__()
100
+ self.num_groups = num_groups
101
+ self.eps = eps
102
+ self.cond_key = cond_key
103
+ self.mapper = nn.Linear(feats_in, c_out * 2)
104
+
105
+ def forward(self, input, cond):
106
+ weight, bias = self.mapper(cond[self.cond_key]).chunk(2, dim=-1)
107
+ input = F.group_norm(input, self.num_groups, eps=self.eps)
108
+ return torch.addcmul(utils.append_dims(bias, input.ndim), input, utils.append_dims(weight, input.ndim) + 1)
109
+
110
+
111
+ # Attention
112
+
113
+ class SelfAttention2d(ConditionedModule):
114
+ def __init__(self, c_in, n_head, norm, dropout_rate=0.):
115
+ super().__init__()
116
+ assert c_in % n_head == 0
117
+ self.norm_in = norm(c_in)
118
+ self.n_head = n_head
119
+ self.qkv_proj = nn.Conv2d(c_in, c_in * 3, 1)
120
+ self.out_proj = nn.Conv2d(c_in, c_in, 1)
121
+ self.dropout = nn.Dropout(dropout_rate)
122
+
123
+ def forward(self, input, cond):
124
+ n, c, h, w = input.shape
125
+ qkv = self.qkv_proj(self.norm_in(input, cond))
126
+ qkv = qkv.view([n, self.n_head * 3, c // self.n_head, h * w]).transpose(2, 3)
127
+ q, k, v = qkv.chunk(3, dim=1)
128
+ scale = k.shape[3] ** -0.25
129
+ att = ((q * scale) @ (k.transpose(2, 3) * scale)).softmax(3)
130
+ att = self.dropout(att)
131
+ y = (att @ v).transpose(2, 3).contiguous().view([n, c, h, w])
132
+ return input + self.out_proj(y)
133
+
134
+
135
+ class CrossAttention2d(ConditionedModule):
136
+ def __init__(self, c_dec, c_enc, n_head, norm_dec, dropout_rate=0.,
137
+ cond_key='cross', cond_key_padding='cross_padding'):
138
+ super().__init__()
139
+ assert c_dec % n_head == 0
140
+ self.cond_key = cond_key
141
+ self.cond_key_padding = cond_key_padding
142
+ self.norm_enc = nn.LayerNorm(c_enc)
143
+ self.norm_dec = norm_dec(c_dec)
144
+ self.n_head = n_head
145
+ self.q_proj = nn.Conv2d(c_dec, c_dec, 1)
146
+ self.kv_proj = nn.Linear(c_enc, c_dec * 2)
147
+ self.out_proj = nn.Conv2d(c_dec, c_dec, 1)
148
+ self.dropout = nn.Dropout(dropout_rate)
149
+
150
+ def forward(self, input, cond):
151
+ n, c, h, w = input.shape
152
+ q = self.q_proj(self.norm_dec(input, cond))
153
+ q = q.view([n, self.n_head, c // self.n_head, h * w]).transpose(2, 3)
154
+ kv = self.kv_proj(self.norm_enc(cond[self.cond_key]))
155
+ kv = kv.view([n, -1, self.n_head * 2, c // self.n_head]).transpose(1, 2)
156
+ k, v = kv.chunk(2, dim=1)
157
+ scale = k.shape[3] ** -0.25
158
+ att = ((q * scale) @ (k.transpose(2, 3) * scale))
159
+ att = att - (cond[self.cond_key_padding][:, None, None, :]) * 10000
160
+ att = att.softmax(3)
161
+ att = self.dropout(att)
162
+ y = (att @ v).transpose(2, 3)
163
+ y = y.contiguous().view([n, c, h, w])
164
+ return input + self.out_proj(y)
165
+
166
+
167
+ # Downsampling/upsampling
168
+
169
+ _kernels = {
170
+ 'linear':
171
+ [1 / 8, 3 / 8, 3 / 8, 1 / 8],
172
+ 'cubic':
173
+ [-0.01171875, -0.03515625, 0.11328125, 0.43359375,
174
+ 0.43359375, 0.11328125, -0.03515625, -0.01171875],
175
+ 'lanczos3':
176
+ [0.003689131001010537, 0.015056144446134567, -0.03399861603975296,
177
+ -0.066637322306633, 0.13550527393817902, 0.44638532400131226,
178
+ 0.44638532400131226, 0.13550527393817902, -0.066637322306633,
179
+ -0.03399861603975296, 0.015056144446134567, 0.003689131001010537]
180
+ }
181
+ _kernels['bilinear'] = _kernels['linear']
182
+ _kernels['bicubic'] = _kernels['cubic']
183
+
184
+
185
+ class Downsample2d(nn.Module):
186
+ def __init__(self, kernel='linear', pad_mode='reflect'):
187
+ super().__init__()
188
+ self.pad_mode = pad_mode
189
+ kernel_1d = torch.tensor([_kernels[kernel]])
190
+ self.pad = kernel_1d.shape[1] // 2 - 1
191
+ self.register_buffer('kernel', kernel_1d.T @ kernel_1d)
192
+
193
+ def forward(self, x):
194
+ x = F.pad(x, (self.pad,) * 4, self.pad_mode)
195
+ weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0], self.kernel.shape[1]])
196
+ indices = torch.arange(x.shape[1], device=x.device)
197
+ weight[indices, indices] = self.kernel.to(weight)
198
+ return F.conv2d(x, weight, stride=2)
199
+
200
+
201
+ class Upsample2d(nn.Module):
202
+ def __init__(self, kernel='linear', pad_mode='reflect'):
203
+ super().__init__()
204
+ self.pad_mode = pad_mode
205
+ kernel_1d = torch.tensor([_kernels[kernel]]) * 2
206
+ self.pad = kernel_1d.shape[1] // 2 - 1
207
+ self.register_buffer('kernel', kernel_1d.T @ kernel_1d)
208
+
209
+ def forward(self, x):
210
+ x = F.pad(x, ((self.pad + 1) // 2,) * 4, self.pad_mode)
211
+ weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0], self.kernel.shape[1]])
212
+ indices = torch.arange(x.shape[1], device=x.device)
213
+ weight[indices, indices] = self.kernel.to(weight)
214
+ return F.conv_transpose2d(x, weight, stride=2, padding=self.pad * 2 + 1)
215
+
216
+
217
+ # Embeddings
218
+
219
+ class FourierFeatures(nn.Module):
220
+ def __init__(self, in_features, out_features, std=1.):
221
+ super().__init__()
222
+ assert out_features % 2 == 0
223
+ self.register_buffer('weight', torch.randn([out_features // 2, in_features]) * std)
224
+
225
+ def forward(self, input):
226
+ f = 2 * math.pi * input @ self.weight.T
227
+ return torch.cat([f.cos(), f.sin()], dim=-1)
228
+
229
+
230
+ # U-Nets
231
+
232
+ class UNet(ConditionedModule):
233
+ def __init__(self, d_blocks, u_blocks, skip_stages=0):
234
+ super().__init__()
235
+ self.d_blocks = nn.ModuleList(d_blocks)
236
+ self.u_blocks = nn.ModuleList(u_blocks)
237
+ self.skip_stages = skip_stages
238
+
239
+ def forward(self, input, cond):
240
+ skips = []
241
+ for block in self.d_blocks[self.skip_stages:]:
242
+ input = block(input, cond)
243
+ skips.append(input)
244
+ for i, (block, skip) in enumerate(zip(self.u_blocks, reversed(skips))):
245
+ input = block(input, cond, skip if i > 0 else None)
246
+ return input
comfy/k_diffusion/models/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .image_v1 import ImageDenoiserModelV1
comfy/k_diffusion/models/image_v1.py ADDED
@@ -0,0 +1,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn import functional as F
6
+
7
+ from .. import layers, utils
8
+
9
+
10
+ def orthogonal_(module):
11
+ nn.init.orthogonal_(module.weight)
12
+ return module
13
+
14
+
15
+ class ResConvBlock(layers.ConditionedResidualBlock):
16
+ def __init__(self, feats_in, c_in, c_mid, c_out, group_size=32, dropout_rate=0.):
17
+ skip = None if c_in == c_out else orthogonal_(nn.Conv2d(c_in, c_out, 1, bias=False))
18
+ super().__init__(
19
+ layers.AdaGN(feats_in, c_in, max(1, c_in // group_size)),
20
+ nn.GELU(),
21
+ nn.Conv2d(c_in, c_mid, 3, padding=1),
22
+ nn.Dropout2d(dropout_rate, inplace=True),
23
+ layers.AdaGN(feats_in, c_mid, max(1, c_mid // group_size)),
24
+ nn.GELU(),
25
+ nn.Conv2d(c_mid, c_out, 3, padding=1),
26
+ nn.Dropout2d(dropout_rate, inplace=True),
27
+ skip=skip)
28
+
29
+
30
+ class DBlock(layers.ConditionedSequential):
31
+ def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., downsample=False, self_attn=False, cross_attn=False, c_enc=0):
32
+ modules = [nn.Identity()]
33
+ for i in range(n_layers):
34
+ my_c_in = c_in if i == 0 else c_mid
35
+ my_c_out = c_mid if i < n_layers - 1 else c_out
36
+ modules.append(ResConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
37
+ if self_attn:
38
+ norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
39
+ modules.append(layers.SelfAttention2d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
40
+ if cross_attn:
41
+ norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
42
+ modules.append(layers.CrossAttention2d(my_c_out, c_enc, max(1, my_c_out // head_size), norm, dropout_rate))
43
+ super().__init__(*modules)
44
+ self.set_downsample(downsample)
45
+
46
+ def set_downsample(self, downsample):
47
+ self[0] = layers.Downsample2d() if downsample else nn.Identity()
48
+ return self
49
+
50
+
51
+ class UBlock(layers.ConditionedSequential):
52
+ def __init__(self, n_layers, feats_in, c_in, c_mid, c_out, group_size=32, head_size=64, dropout_rate=0., upsample=False, self_attn=False, cross_attn=False, c_enc=0):
53
+ modules = []
54
+ for i in range(n_layers):
55
+ my_c_in = c_in if i == 0 else c_mid
56
+ my_c_out = c_mid if i < n_layers - 1 else c_out
57
+ modules.append(ResConvBlock(feats_in, my_c_in, c_mid, my_c_out, group_size, dropout_rate))
58
+ if self_attn:
59
+ norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
60
+ modules.append(layers.SelfAttention2d(my_c_out, max(1, my_c_out // head_size), norm, dropout_rate))
61
+ if cross_attn:
62
+ norm = lambda c_in: layers.AdaGN(feats_in, c_in, max(1, my_c_out // group_size))
63
+ modules.append(layers.CrossAttention2d(my_c_out, c_enc, max(1, my_c_out // head_size), norm, dropout_rate))
64
+ modules.append(nn.Identity())
65
+ super().__init__(*modules)
66
+ self.set_upsample(upsample)
67
+
68
+ def forward(self, input, cond, skip=None):
69
+ if skip is not None:
70
+ input = torch.cat([input, skip], dim=1)
71
+ return super().forward(input, cond)
72
+
73
+ def set_upsample(self, upsample):
74
+ self[-1] = layers.Upsample2d() if upsample else nn.Identity()
75
+ return self
76
+
77
+
78
+ class MappingNet(nn.Sequential):
79
+ def __init__(self, feats_in, feats_out, n_layers=2):
80
+ layers = []
81
+ for i in range(n_layers):
82
+ layers.append(orthogonal_(nn.Linear(feats_in if i == 0 else feats_out, feats_out)))
83
+ layers.append(nn.GELU())
84
+ super().__init__(*layers)
85
+
86
+
87
+ class ImageDenoiserModelV1(nn.Module):
88
+ def __init__(self, c_in, feats_in, depths, channels, self_attn_depths, cross_attn_depths=None, mapping_cond_dim=0, unet_cond_dim=0, cross_cond_dim=0, dropout_rate=0., patch_size=1, skip_stages=0, has_variance=False):
89
+ super().__init__()
90
+ self.c_in = c_in
91
+ self.channels = channels
92
+ self.unet_cond_dim = unet_cond_dim
93
+ self.patch_size = patch_size
94
+ self.has_variance = has_variance
95
+ self.timestep_embed = layers.FourierFeatures(1, feats_in)
96
+ if mapping_cond_dim > 0:
97
+ self.mapping_cond = nn.Linear(mapping_cond_dim, feats_in, bias=False)
98
+ self.mapping = MappingNet(feats_in, feats_in)
99
+ self.proj_in = nn.Conv2d((c_in + unet_cond_dim) * self.patch_size ** 2, channels[max(0, skip_stages - 1)], 1)
100
+ self.proj_out = nn.Conv2d(channels[max(0, skip_stages - 1)], c_in * self.patch_size ** 2 + (1 if self.has_variance else 0), 1)
101
+ nn.init.zeros_(self.proj_out.weight)
102
+ nn.init.zeros_(self.proj_out.bias)
103
+ if cross_cond_dim == 0:
104
+ cross_attn_depths = [False] * len(self_attn_depths)
105
+ d_blocks, u_blocks = [], []
106
+ for i in range(len(depths)):
107
+ my_c_in = channels[max(0, i - 1)]
108
+ d_blocks.append(DBlock(depths[i], feats_in, my_c_in, channels[i], channels[i], downsample=i > skip_stages, self_attn=self_attn_depths[i], cross_attn=cross_attn_depths[i], c_enc=cross_cond_dim, dropout_rate=dropout_rate))
109
+ for i in range(len(depths)):
110
+ my_c_in = channels[i] * 2 if i < len(depths) - 1 else channels[i]
111
+ my_c_out = channels[max(0, i - 1)]
112
+ u_blocks.append(UBlock(depths[i], feats_in, my_c_in, channels[i], my_c_out, upsample=i > skip_stages, self_attn=self_attn_depths[i], cross_attn=cross_attn_depths[i], c_enc=cross_cond_dim, dropout_rate=dropout_rate))
113
+ self.u_net = layers.UNet(d_blocks, reversed(u_blocks), skip_stages=skip_stages)
114
+
115
+ def forward(self, input, sigma, mapping_cond=None, unet_cond=None, cross_cond=None, cross_cond_padding=None, return_variance=False):
116
+ c_noise = sigma.log() / 4
117
+ timestep_embed = self.timestep_embed(utils.append_dims(c_noise, 2))
118
+ mapping_cond_embed = torch.zeros_like(timestep_embed) if mapping_cond is None else self.mapping_cond(mapping_cond)
119
+ mapping_out = self.mapping(timestep_embed + mapping_cond_embed)
120
+ cond = {'cond': mapping_out}
121
+ if unet_cond is not None:
122
+ input = torch.cat([input, unet_cond], dim=1)
123
+ if cross_cond is not None:
124
+ cond['cross'] = cross_cond
125
+ cond['cross_padding'] = cross_cond_padding
126
+ if self.patch_size > 1:
127
+ input = F.pixel_unshuffle(input, self.patch_size)
128
+ input = self.proj_in(input)
129
+ input = self.u_net(input, cond)
130
+ input = self.proj_out(input)
131
+ if self.has_variance:
132
+ input, logvar = input[:, :-1], input[:, -1].flatten(1).mean(1)
133
+ if self.patch_size > 1:
134
+ input = F.pixel_shuffle(input, self.patch_size)
135
+ if self.has_variance and return_variance:
136
+ return input, logvar
137
+ return input
138
+
139
+ def set_skip_stages(self, skip_stages):
140
+ self.proj_in = nn.Conv2d(self.proj_in.in_channels, self.channels[max(0, skip_stages - 1)], 1)
141
+ self.proj_out = nn.Conv2d(self.channels[max(0, skip_stages - 1)], self.proj_out.out_channels, 1)
142
+ nn.init.zeros_(self.proj_out.weight)
143
+ nn.init.zeros_(self.proj_out.bias)
144
+ self.u_net.skip_stages = skip_stages
145
+ for i, block in enumerate(self.u_net.d_blocks):
146
+ block.set_downsample(i > skip_stages)
147
+ for i, block in enumerate(reversed(self.u_net.u_blocks)):
148
+ block.set_upsample(i > skip_stages)
149
+ return self
150
+
151
+ def set_patch_size(self, patch_size):
152
+ self.patch_size = patch_size
153
+ self.proj_in = nn.Conv2d((self.c_in + self.unet_cond_dim) * self.patch_size ** 2, self.channels[max(0, self.u_net.skip_stages - 1)], 1)
154
+ self.proj_out = nn.Conv2d(self.channels[max(0, self.u_net.skip_stages - 1)], self.c_in * self.patch_size ** 2 + (1 if self.has_variance else 0), 1)
155
+ nn.init.zeros_(self.proj_out.weight)
156
+ nn.init.zeros_(self.proj_out.bias)
comfy/k_diffusion/sampling.py ADDED
@@ -0,0 +1,607 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ from scipy import integrate
4
+ import torch
5
+ from torch import nn
6
+ from torchdiffeq import odeint
7
+ import torchsde
8
+ from tqdm.auto import trange, tqdm
9
+
10
+ from . import utils
11
+
12
+
13
+ def append_zero(x):
14
+ return torch.cat([x, x.new_zeros([1])])
15
+
16
+
17
+ def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
18
+ """Constructs the noise schedule of Karras et al. (2022)."""
19
+ ramp = torch.linspace(0, 1, n, device=device)
20
+ min_inv_rho = sigma_min ** (1 / rho)
21
+ max_inv_rho = sigma_max ** (1 / rho)
22
+ sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
23
+ return append_zero(sigmas).to(device)
24
+
25
+
26
+ def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
27
+ """Constructs an exponential noise schedule."""
28
+ sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
29
+ return append_zero(sigmas)
30
+
31
+
32
+ def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
33
+ """Constructs an polynomial in log sigma noise schedule."""
34
+ ramp = torch.linspace(1, 0, n, device=device) ** rho
35
+ sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
36
+ return append_zero(sigmas)
37
+
38
+
39
+ def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
40
+ """Constructs a continuous VP noise schedule."""
41
+ t = torch.linspace(1, eps_s, n, device=device)
42
+ sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
43
+ return append_zero(sigmas)
44
+
45
+
46
+ def to_d(x, sigma, denoised):
47
+ """Converts a denoiser output to a Karras ODE derivative."""
48
+ return (x - denoised) / utils.append_dims(sigma, x.ndim)
49
+
50
+
51
+ def get_ancestral_step(sigma_from, sigma_to, eta=1.):
52
+ """Calculates the noise level (sigma_down) to step down to and the amount
53
+ of noise to add (sigma_up) when doing an ancestral sampling step."""
54
+ if not eta:
55
+ return sigma_to, 0.
56
+ sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
57
+ sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
58
+ return sigma_down, sigma_up
59
+
60
+
61
+ def default_noise_sampler(x):
62
+ return lambda sigma, sigma_next: torch.randn_like(x)
63
+
64
+
65
+ class BatchedBrownianTree:
66
+ """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
67
+
68
+ def __init__(self, x, t0, t1, seed=None, **kwargs):
69
+ t0, t1, self.sign = self.sort(t0, t1)
70
+ w0 = kwargs.get('w0', torch.zeros_like(x))
71
+ if seed is None:
72
+ seed = torch.randint(0, 2 ** 63 - 1, []).item()
73
+ self.batched = True
74
+ try:
75
+ assert len(seed) == x.shape[0]
76
+ w0 = w0[0]
77
+ except TypeError:
78
+ seed = [seed]
79
+ self.batched = False
80
+ self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
81
+
82
+ @staticmethod
83
+ def sort(a, b):
84
+ return (a, b, 1) if a < b else (b, a, -1)
85
+
86
+ def __call__(self, t0, t1):
87
+ t0, t1, sign = self.sort(t0, t1)
88
+ w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
89
+ return w if self.batched else w[0]
90
+
91
+
92
+ class BrownianTreeNoiseSampler:
93
+ """A noise sampler backed by a torchsde.BrownianTree.
94
+
95
+ Args:
96
+ x (Tensor): The tensor whose shape, device and dtype to use to generate
97
+ random samples.
98
+ sigma_min (float): The low end of the valid interval.
99
+ sigma_max (float): The high end of the valid interval.
100
+ seed (int or List[int]): The random seed. If a list of seeds is
101
+ supplied instead of a single integer, then the noise sampler will
102
+ use one BrownianTree per batch item, each with its own seed.
103
+ transform (callable): A function that maps sigma to the sampler's
104
+ internal timestep.
105
+ """
106
+
107
+ def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
108
+ self.transform = transform
109
+ t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
110
+ self.tree = BatchedBrownianTree(x, t0, t1, seed)
111
+
112
+ def __call__(self, sigma, sigma_next):
113
+ t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
114
+ return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
115
+
116
+
117
+ @torch.no_grad()
118
+ def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
119
+ """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
120
+ extra_args = {} if extra_args is None else extra_args
121
+ s_in = x.new_ones([x.shape[0]])
122
+ for i in trange(len(sigmas) - 1, disable=disable):
123
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
124
+ eps = torch.randn_like(x) * s_noise
125
+ sigma_hat = sigmas[i] * (gamma + 1)
126
+ if gamma > 0:
127
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
128
+ denoised = model(x, sigma_hat * s_in, **extra_args)
129
+ d = to_d(x, sigma_hat, denoised)
130
+ if callback is not None:
131
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
132
+ dt = sigmas[i + 1] - sigma_hat
133
+ # Euler method
134
+ x = x + d * dt
135
+ return x
136
+
137
+
138
+ @torch.no_grad()
139
+ def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
140
+ """Ancestral sampling with Euler method steps."""
141
+ extra_args = {} if extra_args is None else extra_args
142
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
143
+ s_in = x.new_ones([x.shape[0]])
144
+ for i in trange(len(sigmas) - 1, disable=disable):
145
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
146
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
147
+ if callback is not None:
148
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
149
+ d = to_d(x, sigmas[i], denoised)
150
+ # Euler method
151
+ dt = sigma_down - sigmas[i]
152
+ x = x + d * dt
153
+ if sigmas[i + 1] > 0:
154
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
155
+ return x
156
+
157
+
158
+ @torch.no_grad()
159
+ def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
160
+ """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
161
+ extra_args = {} if extra_args is None else extra_args
162
+ s_in = x.new_ones([x.shape[0]])
163
+ for i in trange(len(sigmas) - 1, disable=disable):
164
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
165
+ eps = torch.randn_like(x) * s_noise
166
+ sigma_hat = sigmas[i] * (gamma + 1)
167
+ if gamma > 0:
168
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
169
+ denoised = model(x, sigma_hat * s_in, **extra_args)
170
+ d = to_d(x, sigma_hat, denoised)
171
+ if callback is not None:
172
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
173
+ dt = sigmas[i + 1] - sigma_hat
174
+ if sigmas[i + 1] == 0:
175
+ # Euler method
176
+ x = x + d * dt
177
+ else:
178
+ # Heun's method
179
+ x_2 = x + d * dt
180
+ denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
181
+ d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
182
+ d_prime = (d + d_2) / 2
183
+ x = x + d_prime * dt
184
+ return x
185
+
186
+
187
+ @torch.no_grad()
188
+ def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
189
+ """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
190
+ extra_args = {} if extra_args is None else extra_args
191
+ s_in = x.new_ones([x.shape[0]])
192
+ for i in trange(len(sigmas) - 1, disable=disable):
193
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
194
+ eps = torch.randn_like(x) * s_noise
195
+ sigma_hat = sigmas[i] * (gamma + 1)
196
+ if gamma > 0:
197
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
198
+ denoised = model(x, sigma_hat * s_in, **extra_args)
199
+ d = to_d(x, sigma_hat, denoised)
200
+ if callback is not None:
201
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
202
+ if sigmas[i + 1] == 0:
203
+ # Euler method
204
+ dt = sigmas[i + 1] - sigma_hat
205
+ x = x + d * dt
206
+ else:
207
+ # DPM-Solver-2
208
+ sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
209
+ dt_1 = sigma_mid - sigma_hat
210
+ dt_2 = sigmas[i + 1] - sigma_hat
211
+ x_2 = x + d * dt_1
212
+ denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
213
+ d_2 = to_d(x_2, sigma_mid, denoised_2)
214
+ x = x + d_2 * dt_2
215
+ return x
216
+
217
+
218
+ @torch.no_grad()
219
+ def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
220
+ """Ancestral sampling with DPM-Solver second-order steps."""
221
+ extra_args = {} if extra_args is None else extra_args
222
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
223
+ s_in = x.new_ones([x.shape[0]])
224
+ for i in trange(len(sigmas) - 1, disable=disable):
225
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
226
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
227
+ if callback is not None:
228
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
229
+ d = to_d(x, sigmas[i], denoised)
230
+ if sigma_down == 0:
231
+ # Euler method
232
+ dt = sigma_down - sigmas[i]
233
+ x = x + d * dt
234
+ else:
235
+ # DPM-Solver-2
236
+ sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
237
+ dt_1 = sigma_mid - sigmas[i]
238
+ dt_2 = sigma_down - sigmas[i]
239
+ x_2 = x + d * dt_1
240
+ denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
241
+ d_2 = to_d(x_2, sigma_mid, denoised_2)
242
+ x = x + d_2 * dt_2
243
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
244
+ return x
245
+
246
+
247
+ def linear_multistep_coeff(order, t, i, j):
248
+ if order - 1 > i:
249
+ raise ValueError(f'Order {order} too high for step {i}')
250
+ def fn(tau):
251
+ prod = 1.
252
+ for k in range(order):
253
+ if j == k:
254
+ continue
255
+ prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
256
+ return prod
257
+ return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
258
+
259
+
260
+ @torch.no_grad()
261
+ def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
262
+ extra_args = {} if extra_args is None else extra_args
263
+ s_in = x.new_ones([x.shape[0]])
264
+ sigmas_cpu = sigmas.detach().cpu().numpy()
265
+ ds = []
266
+ for i in trange(len(sigmas) - 1, disable=disable):
267
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
268
+ d = to_d(x, sigmas[i], denoised)
269
+ ds.append(d)
270
+ if len(ds) > order:
271
+ ds.pop(0)
272
+ if callback is not None:
273
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
274
+ cur_order = min(i + 1, order)
275
+ coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
276
+ x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
277
+ return x
278
+
279
+
280
+ @torch.no_grad()
281
+ def log_likelihood(model, x, sigma_min, sigma_max, extra_args=None, atol=1e-4, rtol=1e-4):
282
+ extra_args = {} if extra_args is None else extra_args
283
+ s_in = x.new_ones([x.shape[0]])
284
+ v = torch.randint_like(x, 2) * 2 - 1
285
+ fevals = 0
286
+ def ode_fn(sigma, x):
287
+ nonlocal fevals
288
+ with torch.enable_grad():
289
+ x = x[0].detach().requires_grad_()
290
+ denoised = model(x, sigma * s_in, **extra_args)
291
+ d = to_d(x, sigma, denoised)
292
+ fevals += 1
293
+ grad = torch.autograd.grad((d * v).sum(), x)[0]
294
+ d_ll = (v * grad).flatten(1).sum(1)
295
+ return d.detach(), d_ll
296
+ x_min = x, x.new_zeros([x.shape[0]])
297
+ t = x.new_tensor([sigma_min, sigma_max])
298
+ sol = odeint(ode_fn, x_min, t, atol=atol, rtol=rtol, method='dopri5')
299
+ latent, delta_ll = sol[0][-1], sol[1][-1]
300
+ ll_prior = torch.distributions.Normal(0, sigma_max).log_prob(latent).flatten(1).sum(1)
301
+ return ll_prior + delta_ll, {'fevals': fevals}
302
+
303
+
304
+ class PIDStepSizeController:
305
+ """A PID controller for ODE adaptive step size control."""
306
+ def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
307
+ self.h = h
308
+ self.b1 = (pcoeff + icoeff + dcoeff) / order
309
+ self.b2 = -(pcoeff + 2 * dcoeff) / order
310
+ self.b3 = dcoeff / order
311
+ self.accept_safety = accept_safety
312
+ self.eps = eps
313
+ self.errs = []
314
+
315
+ def limiter(self, x):
316
+ return 1 + math.atan(x - 1)
317
+
318
+ def propose_step(self, error):
319
+ inv_error = 1 / (float(error) + self.eps)
320
+ if not self.errs:
321
+ self.errs = [inv_error, inv_error, inv_error]
322
+ self.errs[0] = inv_error
323
+ factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
324
+ factor = self.limiter(factor)
325
+ accept = factor >= self.accept_safety
326
+ if accept:
327
+ self.errs[2] = self.errs[1]
328
+ self.errs[1] = self.errs[0]
329
+ self.h *= factor
330
+ return accept
331
+
332
+
333
+ class DPMSolver(nn.Module):
334
+ """DPM-Solver. See https://arxiv.org/abs/2206.00927."""
335
+
336
+ def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None):
337
+ super().__init__()
338
+ self.model = model
339
+ self.extra_args = {} if extra_args is None else extra_args
340
+ self.eps_callback = eps_callback
341
+ self.info_callback = info_callback
342
+
343
+ def t(self, sigma):
344
+ return -sigma.log()
345
+
346
+ def sigma(self, t):
347
+ return t.neg().exp()
348
+
349
+ def eps(self, eps_cache, key, x, t, *args, **kwargs):
350
+ if key in eps_cache:
351
+ return eps_cache[key], eps_cache
352
+ sigma = self.sigma(t) * x.new_ones([x.shape[0]])
353
+ eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t)
354
+ if self.eps_callback is not None:
355
+ self.eps_callback()
356
+ return eps, {key: eps, **eps_cache}
357
+
358
+ def dpm_solver_1_step(self, x, t, t_next, eps_cache=None):
359
+ eps_cache = {} if eps_cache is None else eps_cache
360
+ h = t_next - t
361
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
362
+ x_1 = x - self.sigma(t_next) * h.expm1() * eps
363
+ return x_1, eps_cache
364
+
365
+ def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None):
366
+ eps_cache = {} if eps_cache is None else eps_cache
367
+ h = t_next - t
368
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
369
+ s1 = t + r1 * h
370
+ u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
371
+ eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
372
+ x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
373
+ return x_2, eps_cache
374
+
375
+ def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None):
376
+ eps_cache = {} if eps_cache is None else eps_cache
377
+ h = t_next - t
378
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
379
+ s1 = t + r1 * h
380
+ s2 = t + r2 * h
381
+ u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
382
+ eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
383
+ u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
384
+ eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2)
385
+ x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
386
+ return x_3, eps_cache
387
+
388
+ def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
389
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
390
+ if not t_end > t_start and eta:
391
+ raise ValueError('eta must be 0 for reverse sampling')
392
+
393
+ m = math.floor(nfe / 3) + 1
394
+ ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
395
+
396
+ if nfe % 3 == 0:
397
+ orders = [3] * (m - 2) + [2, 1]
398
+ else:
399
+ orders = [3] * (m - 1) + [nfe % 3]
400
+
401
+ for i in range(len(orders)):
402
+ eps_cache = {}
403
+ t, t_next = ts[i], ts[i + 1]
404
+ if eta:
405
+ sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
406
+ t_next_ = torch.minimum(t_end, self.t(sd))
407
+ su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
408
+ else:
409
+ t_next_, su = t_next, 0.
410
+
411
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
412
+ denoised = x - self.sigma(t) * eps
413
+ if self.info_callback is not None:
414
+ self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
415
+
416
+ if orders[i] == 1:
417
+ x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache)
418
+ elif orders[i] == 2:
419
+ x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)
420
+ else:
421
+ x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)
422
+
423
+ x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
424
+
425
+ return x
426
+
427
+ def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
428
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
429
+ if order not in {2, 3}:
430
+ raise ValueError('order should be 2 or 3')
431
+ forward = t_end > t_start
432
+ if not forward and eta:
433
+ raise ValueError('eta must be 0 for reverse sampling')
434
+ h_init = abs(h_init) * (1 if forward else -1)
435
+ atol = torch.tensor(atol)
436
+ rtol = torch.tensor(rtol)
437
+ s = t_start
438
+ x_prev = x
439
+ accept = True
440
+ pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
441
+ info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
442
+
443
+ while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
444
+ eps_cache = {}
445
+ t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
446
+ if eta:
447
+ sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
448
+ t_ = torch.minimum(t_end, self.t(sd))
449
+ su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
450
+ else:
451
+ t_, su = t, 0.
452
+
453
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
454
+ denoised = x - self.sigma(s) * eps
455
+
456
+ if order == 2:
457
+ x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
458
+ x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
459
+ else:
460
+ x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
461
+ x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
462
+ delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
463
+ error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
464
+ accept = pid.propose_step(error)
465
+ if accept:
466
+ x_prev = x_low
467
+ x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
468
+ s = t
469
+ info['n_accept'] += 1
470
+ else:
471
+ info['n_reject'] += 1
472
+ info['nfe'] += order
473
+ info['steps'] += 1
474
+
475
+ if self.info_callback is not None:
476
+ self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
477
+
478
+ return x, info
479
+
480
+
481
+ @torch.no_grad()
482
+ def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None):
483
+ """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
484
+ if sigma_min <= 0 or sigma_max <= 0:
485
+ raise ValueError('sigma_min and sigma_max must not be 0')
486
+ with tqdm(total=n, disable=disable) as pbar:
487
+ dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
488
+ if callback is not None:
489
+ dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
490
+ return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), n, eta, s_noise, noise_sampler)
491
+
492
+
493
+ @torch.no_grad()
494
+ def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callback=None, disable=None, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None, return_info=False):
495
+ """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
496
+ if sigma_min <= 0 or sigma_max <= 0:
497
+ raise ValueError('sigma_min and sigma_max must not be 0')
498
+ with tqdm(disable=disable) as pbar:
499
+ dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
500
+ if callback is not None:
501
+ dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
502
+ x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
503
+ if return_info:
504
+ return x, info
505
+ return x
506
+
507
+
508
+ @torch.no_grad()
509
+ def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
510
+ """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
511
+ extra_args = {} if extra_args is None else extra_args
512
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
513
+ s_in = x.new_ones([x.shape[0]])
514
+ sigma_fn = lambda t: t.neg().exp()
515
+ t_fn = lambda sigma: sigma.log().neg()
516
+
517
+ for i in trange(len(sigmas) - 1, disable=disable):
518
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
519
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
520
+ if callback is not None:
521
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
522
+ if sigma_down == 0:
523
+ # Euler method
524
+ d = to_d(x, sigmas[i], denoised)
525
+ dt = sigma_down - sigmas[i]
526
+ x = x + d * dt
527
+ else:
528
+ # DPM-Solver++(2S)
529
+ t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
530
+ r = 1 / 2
531
+ h = t_next - t
532
+ s = t + r * h
533
+ x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
534
+ denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
535
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
536
+ # Noise addition
537
+ if sigmas[i + 1] > 0:
538
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
539
+ return x
540
+
541
+
542
+ @torch.no_grad()
543
+ def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
544
+ """DPM-Solver++ (stochastic)."""
545
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
546
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
547
+ extra_args = {} if extra_args is None else extra_args
548
+ s_in = x.new_ones([x.shape[0]])
549
+ sigma_fn = lambda t: t.neg().exp()
550
+ t_fn = lambda sigma: sigma.log().neg()
551
+
552
+ for i in trange(len(sigmas) - 1, disable=disable):
553
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
554
+ if callback is not None:
555
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
556
+ if sigmas[i + 1] == 0:
557
+ # Euler method
558
+ d = to_d(x, sigmas[i], denoised)
559
+ dt = sigmas[i + 1] - sigmas[i]
560
+ x = x + d * dt
561
+ else:
562
+ # DPM-Solver++
563
+ t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
564
+ h = t_next - t
565
+ s = t + h * r
566
+ fac = 1 / (2 * r)
567
+
568
+ # Step 1
569
+ sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
570
+ s_ = t_fn(sd)
571
+ x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised
572
+ x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
573
+ denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
574
+
575
+ # Step 2
576
+ sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
577
+ t_next_ = t_fn(sd)
578
+ denoised_d = (1 - fac) * denoised + fac * denoised_2
579
+ x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d
580
+ x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
581
+ return x
582
+
583
+
584
+ @torch.no_grad()
585
+ def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
586
+ """DPM-Solver++(2M)."""
587
+ extra_args = {} if extra_args is None else extra_args
588
+ s_in = x.new_ones([x.shape[0]])
589
+ sigma_fn = lambda t: t.neg().exp()
590
+ t_fn = lambda sigma: sigma.log().neg()
591
+ old_denoised = None
592
+
593
+ for i in trange(len(sigmas) - 1, disable=disable):
594
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
595
+ if callback is not None:
596
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
597
+ t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
598
+ h = t_next - t
599
+ if old_denoised is None or sigmas[i + 1] == 0:
600
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
601
+ else:
602
+ h_last = t - t_fn(sigmas[i - 1])
603
+ r = h_last / h
604
+ denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
605
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
606
+ old_denoised = denoised
607
+ return x
comfy/k_diffusion/utils.py ADDED
@@ -0,0 +1,332 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from contextlib import contextmanager
2
+ import hashlib
3
+ import math
4
+ from pathlib import Path
5
+ import shutil
6
+ import urllib
7
+ import warnings
8
+
9
+ from PIL import Image
10
+ import torch
11
+ from torch import nn, optim
12
+ from torch.utils import data
13
+ from torchvision.transforms import functional as TF
14
+
15
+
16
+ def from_pil_image(x):
17
+ """Converts from a PIL image to a tensor."""
18
+ x = TF.to_tensor(x)
19
+ if x.ndim == 2:
20
+ x = x[..., None]
21
+ return x * 2 - 1
22
+
23
+
24
+ def to_pil_image(x):
25
+ """Converts from a tensor to a PIL image."""
26
+ if x.ndim == 4:
27
+ assert x.shape[0] == 1
28
+ x = x[0]
29
+ if x.shape[0] == 1:
30
+ x = x[0]
31
+ return TF.to_pil_image((x.clamp(-1, 1) + 1) / 2)
32
+
33
+
34
+ def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'):
35
+ """Apply passed in transforms for HuggingFace Datasets."""
36
+ images = [transform(image.convert(mode)) for image in examples[image_key]]
37
+ return {image_key: images}
38
+
39
+
40
+ def append_dims(x, target_dims):
41
+ """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
42
+ dims_to_append = target_dims - x.ndim
43
+ if dims_to_append < 0:
44
+ raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
45
+ expanded = x[(...,) + (None,) * dims_to_append]
46
+ # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
47
+ # https://github.com/pytorch/pytorch/issues/84364
48
+ return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
49
+
50
+
51
+ def n_params(module):
52
+ """Returns the number of trainable parameters in a module."""
53
+ return sum(p.numel() for p in module.parameters())
54
+
55
+
56
+ def download_file(path, url, digest=None):
57
+ """Downloads a file if it does not exist, optionally checking its SHA-256 hash."""
58
+ path = Path(path)
59
+ path.parent.mkdir(parents=True, exist_ok=True)
60
+ if not path.exists():
61
+ with urllib.request.urlopen(url) as response, open(path, 'wb') as f:
62
+ shutil.copyfileobj(response, f)
63
+ if digest is not None:
64
+ file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()
65
+ if digest != file_digest:
66
+ raise OSError(f'hash of {path} (url: {url}) failed to validate')
67
+ return path
68
+
69
+
70
+ @contextmanager
71
+ def train_mode(model, mode=True):
72
+ """A context manager that places a model into training mode and restores
73
+ the previous mode on exit."""
74
+ modes = [module.training for module in model.modules()]
75
+ try:
76
+ yield model.train(mode)
77
+ finally:
78
+ for i, module in enumerate(model.modules()):
79
+ module.training = modes[i]
80
+
81
+
82
+ def eval_mode(model):
83
+ """A context manager that places a model into evaluation mode and restores
84
+ the previous mode on exit."""
85
+ return train_mode(model, False)
86
+
87
+
88
+ @torch.no_grad()
89
+ def ema_update(model, averaged_model, decay):
90
+ """Incorporates updated model parameters into an exponential moving averaged
91
+ version of a model. It should be called after each optimizer step."""
92
+ model_params = dict(model.named_parameters())
93
+ averaged_params = dict(averaged_model.named_parameters())
94
+ assert model_params.keys() == averaged_params.keys()
95
+
96
+ for name, param in model_params.items():
97
+ averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)
98
+
99
+ model_buffers = dict(model.named_buffers())
100
+ averaged_buffers = dict(averaged_model.named_buffers())
101
+ assert model_buffers.keys() == averaged_buffers.keys()
102
+
103
+ for name, buf in model_buffers.items():
104
+ averaged_buffers[name].copy_(buf)
105
+
106
+
107
+ class EMAWarmup:
108
+ """Implements an EMA warmup using an inverse decay schedule.
109
+ If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
110
+ good values for models you plan to train for a million or more steps (reaches decay
111
+ factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
112
+ you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
113
+ 215.4k steps).
114
+ Args:
115
+ inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
116
+ power (float): Exponential factor of EMA warmup. Default: 1.
117
+ min_value (float): The minimum EMA decay rate. Default: 0.
118
+ max_value (float): The maximum EMA decay rate. Default: 1.
119
+ start_at (int): The epoch to start averaging at. Default: 0.
120
+ last_epoch (int): The index of last epoch. Default: 0.
121
+ """
122
+
123
+ def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
124
+ last_epoch=0):
125
+ self.inv_gamma = inv_gamma
126
+ self.power = power
127
+ self.min_value = min_value
128
+ self.max_value = max_value
129
+ self.start_at = start_at
130
+ self.last_epoch = last_epoch
131
+
132
+ def state_dict(self):
133
+ """Returns the state of the class as a :class:`dict`."""
134
+ return dict(self.__dict__.items())
135
+
136
+ def load_state_dict(self, state_dict):
137
+ """Loads the class's state.
138
+ Args:
139
+ state_dict (dict): scaler state. Should be an object returned
140
+ from a call to :meth:`state_dict`.
141
+ """
142
+ self.__dict__.update(state_dict)
143
+
144
+ def get_value(self):
145
+ """Gets the current EMA decay rate."""
146
+ epoch = max(0, self.last_epoch - self.start_at)
147
+ value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
148
+ return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
149
+
150
+ def step(self):
151
+ """Updates the step count."""
152
+ self.last_epoch += 1
153
+
154
+
155
+ class InverseLR(optim.lr_scheduler._LRScheduler):
156
+ """Implements an inverse decay learning rate schedule with an optional exponential
157
+ warmup. When last_epoch=-1, sets initial lr as lr.
158
+ inv_gamma is the number of steps/epochs required for the learning rate to decay to
159
+ (1 / 2)**power of its original value.
160
+ Args:
161
+ optimizer (Optimizer): Wrapped optimizer.
162
+ inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
163
+ power (float): Exponential factor of learning rate decay. Default: 1.
164
+ warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
165
+ Default: 0.
166
+ min_lr (float): The minimum learning rate. Default: 0.
167
+ last_epoch (int): The index of last epoch. Default: -1.
168
+ verbose (bool): If ``True``, prints a message to stdout for
169
+ each update. Default: ``False``.
170
+ """
171
+
172
+ def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,
173
+ last_epoch=-1, verbose=False):
174
+ self.inv_gamma = inv_gamma
175
+ self.power = power
176
+ if not 0. <= warmup < 1:
177
+ raise ValueError('Invalid value for warmup')
178
+ self.warmup = warmup
179
+ self.min_lr = min_lr
180
+ super().__init__(optimizer, last_epoch, verbose)
181
+
182
+ def get_lr(self):
183
+ if not self._get_lr_called_within_step:
184
+ warnings.warn("To get the last learning rate computed by the scheduler, "
185
+ "please use `get_last_lr()`.")
186
+
187
+ return self._get_closed_form_lr()
188
+
189
+ def _get_closed_form_lr(self):
190
+ warmup = 1 - self.warmup ** (self.last_epoch + 1)
191
+ lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
192
+ return [warmup * max(self.min_lr, base_lr * lr_mult)
193
+ for base_lr in self.base_lrs]
194
+
195
+
196
+ class ExponentialLR(optim.lr_scheduler._LRScheduler):
197
+ """Implements an exponential learning rate schedule with an optional exponential
198
+ warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate
199
+ continuously by decay (default 0.5) every num_steps steps.
200
+ Args:
201
+ optimizer (Optimizer): Wrapped optimizer.
202
+ num_steps (float): The number of steps to decay the learning rate by decay in.
203
+ decay (float): The factor by which to decay the learning rate every num_steps
204
+ steps. Default: 0.5.
205
+ warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
206
+ Default: 0.
207
+ min_lr (float): The minimum learning rate. Default: 0.
208
+ last_epoch (int): The index of last epoch. Default: -1.
209
+ verbose (bool): If ``True``, prints a message to stdout for
210
+ each update. Default: ``False``.
211
+ """
212
+
213
+ def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,
214
+ last_epoch=-1, verbose=False):
215
+ self.num_steps = num_steps
216
+ self.decay = decay
217
+ if not 0. <= warmup < 1:
218
+ raise ValueError('Invalid value for warmup')
219
+ self.warmup = warmup
220
+ self.min_lr = min_lr
221
+ super().__init__(optimizer, last_epoch, verbose)
222
+
223
+ def get_lr(self):
224
+ if not self._get_lr_called_within_step:
225
+ warnings.warn("To get the last learning rate computed by the scheduler, "
226
+ "please use `get_last_lr()`.")
227
+
228
+ return self._get_closed_form_lr()
229
+
230
+ def _get_closed_form_lr(self):
231
+ warmup = 1 - self.warmup ** (self.last_epoch + 1)
232
+ lr_mult = (self.decay ** (1 / self.num_steps)) ** self.last_epoch
233
+ return [warmup * max(self.min_lr, base_lr * lr_mult)
234
+ for base_lr in self.base_lrs]
235
+
236
+
237
+ def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
238
+ """Draws samples from an lognormal distribution."""
239
+ return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()
240
+
241
+
242
+ def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
243
+ """Draws samples from an optionally truncated log-logistic distribution."""
244
+ min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)
245
+ max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)
246
+ min_cdf = min_value.log().sub(loc).div(scale).sigmoid()
247
+ max_cdf = max_value.log().sub(loc).div(scale).sigmoid()
248
+ u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf
249
+ return u.logit().mul(scale).add(loc).exp().to(dtype)
250
+
251
+
252
+ def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):
253
+ """Draws samples from an log-uniform distribution."""
254
+ min_value = math.log(min_value)
255
+ max_value = math.log(max_value)
256
+ return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()
257
+
258
+
259
+ def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
260
+ """Draws samples from a truncated v-diffusion training timestep distribution."""
261
+ min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi
262
+ max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi
263
+ u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf
264
+ return torch.tan(u * math.pi / 2) * sigma_data
265
+
266
+
267
+ def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32):
268
+ """Draws samples from a split lognormal distribution."""
269
+ n = torch.randn(shape, device=device, dtype=dtype).abs()
270
+ u = torch.rand(shape, device=device, dtype=dtype)
271
+ n_left = n * -scale_1 + loc
272
+ n_right = n * scale_2 + loc
273
+ ratio = scale_1 / (scale_1 + scale_2)
274
+ return torch.where(u < ratio, n_left, n_right).exp()
275
+
276
+
277
+ class FolderOfImages(data.Dataset):
278
+ """Recursively finds all images in a directory. It does not support
279
+ classes/targets."""
280
+
281
+ IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'}
282
+
283
+ def __init__(self, root, transform=None):
284
+ super().__init__()
285
+ self.root = Path(root)
286
+ self.transform = nn.Identity() if transform is None else transform
287
+ self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS)
288
+
289
+ def __repr__(self):
290
+ return f'FolderOfImages(root="{self.root}", len: {len(self)})'
291
+
292
+ def __len__(self):
293
+ return len(self.paths)
294
+
295
+ def __getitem__(self, key):
296
+ path = self.paths[key]
297
+ with open(path, 'rb') as f:
298
+ image = Image.open(f).convert('RGB')
299
+ image = self.transform(image)
300
+ return image,
301
+
302
+
303
+ class CSVLogger:
304
+ def __init__(self, filename, columns):
305
+ self.filename = Path(filename)
306
+ self.columns = columns
307
+ if self.filename.exists():
308
+ self.file = open(self.filename, 'a')
309
+ else:
310
+ self.file = open(self.filename, 'w')
311
+ self.write(*self.columns)
312
+
313
+ def write(self, *args):
314
+ print(*args, sep=',', file=self.file, flush=True)
315
+
316
+
317
+ @contextmanager
318
+ def tf32_mode(cudnn=None, matmul=None):
319
+ """A context manager that sets whether TF32 is allowed on cuDNN or matmul."""
320
+ cudnn_old = torch.backends.cudnn.allow_tf32
321
+ matmul_old = torch.backends.cuda.matmul.allow_tf32
322
+ try:
323
+ if cudnn is not None:
324
+ torch.backends.cudnn.allow_tf32 = cudnn
325
+ if matmul is not None:
326
+ torch.backends.cuda.matmul.allow_tf32 = matmul
327
+ yield
328
+ finally:
329
+ if cudnn is not None:
330
+ torch.backends.cudnn.allow_tf32 = cudnn_old
331
+ if matmul is not None:
332
+ torch.backends.cuda.matmul.allow_tf32 = matmul_old
comfy/ldm/data/__init__.py ADDED
File without changes
comfy/ldm/data/util.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ldm.modules.midas.api import load_midas_transform
4
+
5
+
6
+ class AddMiDaS(object):
7
+ def __init__(self, model_type):
8
+ super().__init__()
9
+ self.transform = load_midas_transform(model_type)
10
+
11
+ def pt2np(self, x):
12
+ x = ((x + 1.0) * .5).detach().cpu().numpy()
13
+ return x
14
+
15
+ def np2pt(self, x):
16
+ x = torch.from_numpy(x) * 2 - 1.
17
+ return x
18
+
19
+ def __call__(self, sample):
20
+ # sample['jpg'] is tensor hwc in [-1, 1] at this point
21
+ x = self.pt2np(sample['jpg'])
22
+ x = self.transform({"image": x})["image"]
23
+ sample['midas_in'] = x
24
+ return sample
comfy/ldm/models/autoencoder.py ADDED
@@ -0,0 +1,223 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ # import pytorch_lightning as pl
3
+ import torch.nn.functional as F
4
+ from contextlib import contextmanager
5
+
6
+ from ldm.modules.diffusionmodules.model import Encoder, Decoder
7
+ from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
8
+
9
+ from ldm.util import instantiate_from_config
10
+ from ldm.modules.ema import LitEma
11
+
12
+ # class AutoencoderKL(pl.LightningModule):
13
+ class AutoencoderKL(torch.nn.Module):
14
+ def __init__(self,
15
+ ddconfig,
16
+ lossconfig,
17
+ embed_dim,
18
+ ckpt_path=None,
19
+ ignore_keys=[],
20
+ image_key="image",
21
+ colorize_nlabels=None,
22
+ monitor=None,
23
+ ema_decay=None,
24
+ learn_logvar=False
25
+ ):
26
+ super().__init__()
27
+ self.learn_logvar = learn_logvar
28
+ self.image_key = image_key
29
+ self.encoder = Encoder(**ddconfig)
30
+ self.decoder = Decoder(**ddconfig)
31
+ self.loss = instantiate_from_config(lossconfig)
32
+ assert ddconfig["double_z"]
33
+ self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
34
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
35
+ self.embed_dim = embed_dim
36
+ if colorize_nlabels is not None:
37
+ assert type(colorize_nlabels)==int
38
+ self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
39
+ if monitor is not None:
40
+ self.monitor = monitor
41
+
42
+ self.use_ema = ema_decay is not None
43
+ if self.use_ema:
44
+ self.ema_decay = ema_decay
45
+ assert 0. < ema_decay < 1.
46
+ self.model_ema = LitEma(self, decay=ema_decay)
47
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
48
+
49
+ if ckpt_path is not None:
50
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
51
+
52
+ def init_from_ckpt(self, path, ignore_keys=list()):
53
+ if path.lower().endswith(".safetensors"):
54
+ import safetensors.torch
55
+ sd = safetensors.torch.load_file(path, device="cpu")
56
+ else:
57
+ sd = torch.load(path, map_location="cpu")["state_dict"]
58
+ keys = list(sd.keys())
59
+ for k in keys:
60
+ for ik in ignore_keys:
61
+ if k.startswith(ik):
62
+ print("Deleting key {} from state_dict.".format(k))
63
+ del sd[k]
64
+ self.load_state_dict(sd, strict=False)
65
+ print(f"Restored from {path}")
66
+
67
+ @contextmanager
68
+ def ema_scope(self, context=None):
69
+ if self.use_ema:
70
+ self.model_ema.store(self.parameters())
71
+ self.model_ema.copy_to(self)
72
+ if context is not None:
73
+ print(f"{context}: Switched to EMA weights")
74
+ try:
75
+ yield None
76
+ finally:
77
+ if self.use_ema:
78
+ self.model_ema.restore(self.parameters())
79
+ if context is not None:
80
+ print(f"{context}: Restored training weights")
81
+
82
+ def on_train_batch_end(self, *args, **kwargs):
83
+ if self.use_ema:
84
+ self.model_ema(self)
85
+
86
+ def encode(self, x):
87
+ h = self.encoder(x)
88
+ moments = self.quant_conv(h)
89
+ posterior = DiagonalGaussianDistribution(moments)
90
+ return posterior
91
+
92
+ def decode(self, z):
93
+ z = self.post_quant_conv(z)
94
+ dec = self.decoder(z)
95
+ return dec
96
+
97
+ def forward(self, input, sample_posterior=True):
98
+ posterior = self.encode(input)
99
+ if sample_posterior:
100
+ z = posterior.sample()
101
+ else:
102
+ z = posterior.mode()
103
+ dec = self.decode(z)
104
+ return dec, posterior
105
+
106
+ def get_input(self, batch, k):
107
+ x = batch[k]
108
+ if len(x.shape) == 3:
109
+ x = x[..., None]
110
+ x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
111
+ return x
112
+
113
+ def training_step(self, batch, batch_idx, optimizer_idx):
114
+ inputs = self.get_input(batch, self.image_key)
115
+ reconstructions, posterior = self(inputs)
116
+
117
+ if optimizer_idx == 0:
118
+ # train encoder+decoder+logvar
119
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
120
+ last_layer=self.get_last_layer(), split="train")
121
+ self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
122
+ self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
123
+ return aeloss
124
+
125
+ if optimizer_idx == 1:
126
+ # train the discriminator
127
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
128
+ last_layer=self.get_last_layer(), split="train")
129
+
130
+ self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
131
+ self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
132
+ return discloss
133
+
134
+ def validation_step(self, batch, batch_idx):
135
+ log_dict = self._validation_step(batch, batch_idx)
136
+ with self.ema_scope():
137
+ log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
138
+ return log_dict
139
+
140
+ def _validation_step(self, batch, batch_idx, postfix=""):
141
+ inputs = self.get_input(batch, self.image_key)
142
+ reconstructions, posterior = self(inputs)
143
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
144
+ last_layer=self.get_last_layer(), split="val"+postfix)
145
+
146
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
147
+ last_layer=self.get_last_layer(), split="val"+postfix)
148
+
149
+ self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
150
+ self.log_dict(log_dict_ae)
151
+ self.log_dict(log_dict_disc)
152
+ return self.log_dict
153
+
154
+ def configure_optimizers(self):
155
+ lr = self.learning_rate
156
+ ae_params_list = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(
157
+ self.quant_conv.parameters()) + list(self.post_quant_conv.parameters())
158
+ if self.learn_logvar:
159
+ print(f"{self.__class__.__name__}: Learning logvar")
160
+ ae_params_list.append(self.loss.logvar)
161
+ opt_ae = torch.optim.Adam(ae_params_list,
162
+ lr=lr, betas=(0.5, 0.9))
163
+ opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
164
+ lr=lr, betas=(0.5, 0.9))
165
+ return [opt_ae, opt_disc], []
166
+
167
+ def get_last_layer(self):
168
+ return self.decoder.conv_out.weight
169
+
170
+ @torch.no_grad()
171
+ def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
172
+ log = dict()
173
+ x = self.get_input(batch, self.image_key)
174
+ x = x.to(self.device)
175
+ if not only_inputs:
176
+ xrec, posterior = self(x)
177
+ if x.shape[1] > 3:
178
+ # colorize with random projection
179
+ assert xrec.shape[1] > 3
180
+ x = self.to_rgb(x)
181
+ xrec = self.to_rgb(xrec)
182
+ log["samples"] = self.decode(torch.randn_like(posterior.sample()))
183
+ log["reconstructions"] = xrec
184
+ if log_ema or self.use_ema:
185
+ with self.ema_scope():
186
+ xrec_ema, posterior_ema = self(x)
187
+ if x.shape[1] > 3:
188
+ # colorize with random projection
189
+ assert xrec_ema.shape[1] > 3
190
+ xrec_ema = self.to_rgb(xrec_ema)
191
+ log["samples_ema"] = self.decode(torch.randn_like(posterior_ema.sample()))
192
+ log["reconstructions_ema"] = xrec_ema
193
+ log["inputs"] = x
194
+ return log
195
+
196
+ def to_rgb(self, x):
197
+ assert self.image_key == "segmentation"
198
+ if not hasattr(self, "colorize"):
199
+ self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
200
+ x = F.conv2d(x, weight=self.colorize)
201
+ x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
202
+ return x
203
+
204
+
205
+ class IdentityFirstStage(torch.nn.Module):
206
+ def __init__(self, *args, vq_interface=False, **kwargs):
207
+ self.vq_interface = vq_interface
208
+ super().__init__()
209
+
210
+ def encode(self, x, *args, **kwargs):
211
+ return x
212
+
213
+ def decode(self, x, *args, **kwargs):
214
+ return x
215
+
216
+ def quantize(self, x, *args, **kwargs):
217
+ if self.vq_interface:
218
+ return x, None, [None, None, None]
219
+ return x
220
+
221
+ def forward(self, x, *args, **kwargs):
222
+ return x
223
+
comfy/ldm/models/diffusion/__init__.py ADDED
File without changes
comfy/ldm/models/diffusion/ddim.py ADDED
@@ -0,0 +1,410 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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, extract_into_tensor
8
+
9
+
10
+ class DDIMSampler(object):
11
+ def __init__(self, model, schedule="linear", device=torch.device("cuda"), **kwargs):
12
+ super().__init__()
13
+ self.model = model
14
+ self.ddpm_num_timesteps = model.num_timesteps
15
+ self.schedule = schedule
16
+ self.device = device
17
+
18
+ def register_buffer(self, name, attr):
19
+ if type(attr) == torch.Tensor:
20
+ if attr.device != self.device:
21
+ attr = attr.float().to(self.device)
22
+ setattr(self, name, attr)
23
+
24
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
25
+ 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
+ self.make_schedule_timesteps(ddim_timesteps, ddim_eta=ddim_eta, verbose=verbose)
28
+
29
+ def make_schedule_timesteps(self, ddim_timesteps, ddim_eta=0., verbose=True):
30
+ self.ddim_timesteps = torch.tensor(ddim_timesteps)
31
+ alphas_cumprod = self.model.alphas_cumprod
32
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
33
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.device)
34
+
35
+ self.register_buffer('betas', to_torch(self.model.betas))
36
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
37
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
38
+
39
+ # calculations for diffusion q(x_t | x_{t-1}) and others
40
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
41
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
42
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
43
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
44
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
45
+
46
+ # ddim sampling parameters
47
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
48
+ ddim_timesteps=self.ddim_timesteps,
49
+ eta=ddim_eta,verbose=verbose)
50
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
51
+ self.register_buffer('ddim_alphas', ddim_alphas)
52
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
53
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
54
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
55
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
56
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
57
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
58
+
59
+ @torch.no_grad()
60
+ def sample_custom(self,
61
+ ddim_timesteps,
62
+ conditioning,
63
+ callback=None,
64
+ img_callback=None,
65
+ quantize_x0=False,
66
+ eta=0.,
67
+ mask=None,
68
+ x0=None,
69
+ temperature=1.,
70
+ noise_dropout=0.,
71
+ score_corrector=None,
72
+ corrector_kwargs=None,
73
+ verbose=True,
74
+ x_T=None,
75
+ log_every_t=100,
76
+ unconditional_guidance_scale=1.,
77
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
78
+ dynamic_threshold=None,
79
+ ucg_schedule=None,
80
+ denoise_function=None,
81
+ extra_args=None,
82
+ to_zero=True,
83
+ end_step=None,
84
+ **kwargs
85
+ ):
86
+ self.make_schedule_timesteps(ddim_timesteps=ddim_timesteps, ddim_eta=eta, verbose=verbose)
87
+ samples, intermediates = self.ddim_sampling(conditioning, x_T.shape,
88
+ callback=callback,
89
+ img_callback=img_callback,
90
+ quantize_denoised=quantize_x0,
91
+ mask=mask, x0=x0,
92
+ ddim_use_original_steps=False,
93
+ noise_dropout=noise_dropout,
94
+ temperature=temperature,
95
+ score_corrector=score_corrector,
96
+ corrector_kwargs=corrector_kwargs,
97
+ x_T=x_T,
98
+ log_every_t=log_every_t,
99
+ unconditional_guidance_scale=unconditional_guidance_scale,
100
+ unconditional_conditioning=unconditional_conditioning,
101
+ dynamic_threshold=dynamic_threshold,
102
+ ucg_schedule=ucg_schedule,
103
+ denoise_function=denoise_function,
104
+ extra_args=extra_args,
105
+ to_zero=to_zero,
106
+ end_step=end_step
107
+ )
108
+ return samples, intermediates
109
+
110
+
111
+ @torch.no_grad()
112
+ def sample(self,
113
+ S,
114
+ batch_size,
115
+ shape,
116
+ conditioning=None,
117
+ callback=None,
118
+ normals_sequence=None,
119
+ img_callback=None,
120
+ quantize_x0=False,
121
+ eta=0.,
122
+ mask=None,
123
+ x0=None,
124
+ temperature=1.,
125
+ noise_dropout=0.,
126
+ score_corrector=None,
127
+ corrector_kwargs=None,
128
+ verbose=True,
129
+ x_T=None,
130
+ log_every_t=100,
131
+ unconditional_guidance_scale=1.,
132
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
133
+ dynamic_threshold=None,
134
+ ucg_schedule=None,
135
+ **kwargs
136
+ ):
137
+ if conditioning is not None:
138
+ if isinstance(conditioning, dict):
139
+ ctmp = conditioning[list(conditioning.keys())[0]]
140
+ while isinstance(ctmp, list): ctmp = ctmp[0]
141
+ cbs = ctmp.shape[0]
142
+ if cbs != batch_size:
143
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
144
+
145
+ elif isinstance(conditioning, list):
146
+ for ctmp in conditioning:
147
+ if ctmp.shape[0] != batch_size:
148
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
149
+
150
+ else:
151
+ if conditioning.shape[0] != batch_size:
152
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
153
+
154
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
155
+ # sampling
156
+ C, H, W = shape
157
+ size = (batch_size, C, H, W)
158
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
159
+
160
+ samples, intermediates = self.ddim_sampling(conditioning, size,
161
+ callback=callback,
162
+ img_callback=img_callback,
163
+ quantize_denoised=quantize_x0,
164
+ mask=mask, x0=x0,
165
+ ddim_use_original_steps=False,
166
+ noise_dropout=noise_dropout,
167
+ temperature=temperature,
168
+ score_corrector=score_corrector,
169
+ corrector_kwargs=corrector_kwargs,
170
+ x_T=x_T,
171
+ log_every_t=log_every_t,
172
+ unconditional_guidance_scale=unconditional_guidance_scale,
173
+ unconditional_conditioning=unconditional_conditioning,
174
+ dynamic_threshold=dynamic_threshold,
175
+ ucg_schedule=ucg_schedule,
176
+ denoise_function=None,
177
+ extra_args=None
178
+ )
179
+ return samples, intermediates
180
+
181
+ @torch.no_grad()
182
+ def ddim_sampling(self, cond, shape,
183
+ x_T=None, ddim_use_original_steps=False,
184
+ callback=None, timesteps=None, quantize_denoised=False,
185
+ mask=None, x0=None, img_callback=None, log_every_t=100,
186
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
187
+ unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
188
+ ucg_schedule=None, denoise_function=None, extra_args=None, to_zero=True, end_step=None):
189
+ device = self.model.betas.device
190
+ b = shape[0]
191
+ if x_T is None:
192
+ img = torch.randn(shape, device=device)
193
+ else:
194
+ img = x_T
195
+
196
+ if timesteps is None:
197
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
198
+ elif timesteps is not None and not ddim_use_original_steps:
199
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
200
+ timesteps = self.ddim_timesteps[:subset_end]
201
+
202
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
203
+ time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else timesteps.flip(0)
204
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
205
+ # print(f"Running DDIM Sampling with {total_steps} timesteps")
206
+
207
+ iterator = tqdm(time_range[:end_step], desc='DDIM Sampler', total=end_step)
208
+
209
+ for i, step in enumerate(iterator):
210
+ index = total_steps - i - 1
211
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
212
+
213
+ if mask is not None:
214
+ assert x0 is not None
215
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
216
+ img = img_orig * mask + (1. - mask) * img
217
+
218
+ if ucg_schedule is not None:
219
+ assert len(ucg_schedule) == len(time_range)
220
+ unconditional_guidance_scale = ucg_schedule[i]
221
+
222
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
223
+ quantize_denoised=quantize_denoised, temperature=temperature,
224
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
225
+ corrector_kwargs=corrector_kwargs,
226
+ unconditional_guidance_scale=unconditional_guidance_scale,
227
+ unconditional_conditioning=unconditional_conditioning,
228
+ dynamic_threshold=dynamic_threshold, denoise_function=denoise_function, extra_args=extra_args)
229
+ img, pred_x0 = outs
230
+ if callback: callback(i)
231
+ if img_callback: img_callback(pred_x0, i)
232
+
233
+ if index % log_every_t == 0 or index == total_steps - 1:
234
+ intermediates['x_inter'].append(img)
235
+ intermediates['pred_x0'].append(pred_x0)
236
+
237
+ if to_zero:
238
+ img = pred_x0
239
+ else:
240
+ if ddim_use_original_steps:
241
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
242
+ else:
243
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
244
+ img /= sqrt_alphas_cumprod[index - 1]
245
+
246
+ return img, intermediates
247
+
248
+ @torch.no_grad()
249
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
250
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
251
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
252
+ dynamic_threshold=None, denoise_function=None, extra_args=None):
253
+ b, *_, device = *x.shape, x.device
254
+
255
+ if denoise_function is not None:
256
+ model_output = denoise_function(self.model.apply_model, x, t, **extra_args)
257
+ elif unconditional_conditioning is None or unconditional_guidance_scale == 1.:
258
+ model_output = self.model.apply_model(x, t, c)
259
+ else:
260
+ x_in = torch.cat([x] * 2)
261
+ t_in = torch.cat([t] * 2)
262
+ if isinstance(c, dict):
263
+ assert isinstance(unconditional_conditioning, dict)
264
+ c_in = dict()
265
+ for k in c:
266
+ if isinstance(c[k], list):
267
+ c_in[k] = [torch.cat([
268
+ unconditional_conditioning[k][i],
269
+ c[k][i]]) for i in range(len(c[k]))]
270
+ else:
271
+ c_in[k] = torch.cat([
272
+ unconditional_conditioning[k],
273
+ c[k]])
274
+ elif isinstance(c, list):
275
+ c_in = list()
276
+ assert isinstance(unconditional_conditioning, list)
277
+ for i in range(len(c)):
278
+ c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
279
+ else:
280
+ c_in = torch.cat([unconditional_conditioning, c])
281
+ model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
282
+ model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
283
+
284
+ if self.model.parameterization == "v":
285
+ e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
286
+ else:
287
+ e_t = model_output
288
+
289
+ if score_corrector is not None:
290
+ assert self.model.parameterization == "eps", 'not implemented'
291
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
292
+
293
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
294
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
295
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
296
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
297
+ # select parameters corresponding to the currently considered timestep
298
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
299
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
300
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
301
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
302
+
303
+ # current prediction for x_0
304
+ if self.model.parameterization != "v":
305
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
306
+ else:
307
+ pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
308
+
309
+ if quantize_denoised:
310
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
311
+
312
+ if dynamic_threshold is not None:
313
+ raise NotImplementedError()
314
+
315
+ # direction pointing to x_t
316
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
317
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
318
+ if noise_dropout > 0.:
319
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
320
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
321
+ return x_prev, pred_x0
322
+
323
+ @torch.no_grad()
324
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
325
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None):
326
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
327
+
328
+ assert t_enc <= num_reference_steps
329
+ num_steps = t_enc
330
+
331
+ if use_original_steps:
332
+ alphas_next = self.alphas_cumprod[:num_steps]
333
+ alphas = self.alphas_cumprod_prev[:num_steps]
334
+ else:
335
+ alphas_next = self.ddim_alphas[:num_steps]
336
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
337
+
338
+ x_next = x0
339
+ intermediates = []
340
+ inter_steps = []
341
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
342
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
343
+ if unconditional_guidance_scale == 1.:
344
+ noise_pred = self.model.apply_model(x_next, t, c)
345
+ else:
346
+ assert unconditional_conditioning is not None
347
+ e_t_uncond, noise_pred = torch.chunk(
348
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
349
+ torch.cat((unconditional_conditioning, c))), 2)
350
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
351
+
352
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
353
+ weighted_noise_pred = alphas_next[i].sqrt() * (
354
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
355
+ x_next = xt_weighted + weighted_noise_pred
356
+ if return_intermediates and i % (
357
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
358
+ intermediates.append(x_next)
359
+ inter_steps.append(i)
360
+ elif return_intermediates and i >= num_steps - 2:
361
+ intermediates.append(x_next)
362
+ inter_steps.append(i)
363
+ if callback: callback(i)
364
+
365
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
366
+ if return_intermediates:
367
+ out.update({'intermediates': intermediates})
368
+ return x_next, out
369
+
370
+ @torch.no_grad()
371
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None, max_denoise=False):
372
+ # fast, but does not allow for exact reconstruction
373
+ # t serves as an index to gather the correct alphas
374
+ if use_original_steps:
375
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
376
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
377
+ else:
378
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
379
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
380
+
381
+ if noise is None:
382
+ noise = torch.randn_like(x0)
383
+ if max_denoise:
384
+ noise_multiplier = 1.0
385
+ else:
386
+ noise_multiplier = extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape)
387
+
388
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 + noise_multiplier * noise)
389
+
390
+ @torch.no_grad()
391
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
392
+ use_original_steps=False, callback=None):
393
+
394
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
395
+ timesteps = timesteps[:t_start]
396
+
397
+ time_range = np.flip(timesteps)
398
+ total_steps = timesteps.shape[0]
399
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
400
+
401
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
402
+ x_dec = x_latent
403
+ for i, step in enumerate(iterator):
404
+ index = total_steps - i - 1
405
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
406
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
407
+ unconditional_guidance_scale=unconditional_guidance_scale,
408
+ unconditional_conditioning=unconditional_conditioning)
409
+ if callback: callback(i)
410
+ return x_dec
comfy/ldm/models/diffusion/ddpm.py ADDED
@@ -0,0 +1,1875 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
22
+ from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
23
+ from ldm.modules.ema import LitEma
24
+ from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
25
+ from ldm.models.autoencoder import IdentityFirstStage, AutoencoderKL
26
+ from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
27
+ from ldm.models.diffusion.ddim import DDIMSampler
28
+
29
+
30
+ __conditioning_keys__ = {'concat': 'c_concat',
31
+ 'crossattn': 'c_crossattn',
32
+ 'adm': 'y'}
33
+
34
+
35
+ def disabled_train(self, mode=True):
36
+ """Overwrite model.train with this function to make sure train/eval mode
37
+ does not change anymore."""
38
+ return self
39
+
40
+
41
+ def uniform_on_device(r1, r2, shape, device):
42
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
43
+
44
+ # class DDPM(pl.LightningModule):
45
+ class DDPM(torch.nn.Module):
46
+ # classic DDPM with Gaussian diffusion, in image space
47
+ def __init__(self,
48
+ unet_config,
49
+ timesteps=1000,
50
+ beta_schedule="linear",
51
+ loss_type="l2",
52
+ ckpt_path=None,
53
+ ignore_keys=[],
54
+ load_only_unet=False,
55
+ monitor="val/loss",
56
+ use_ema=True,
57
+ first_stage_key="image",
58
+ image_size=256,
59
+ channels=3,
60
+ log_every_t=100,
61
+ clip_denoised=True,
62
+ linear_start=1e-4,
63
+ linear_end=2e-2,
64
+ cosine_s=8e-3,
65
+ given_betas=None,
66
+ original_elbo_weight=0.,
67
+ v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
68
+ l_simple_weight=1.,
69
+ conditioning_key=None,
70
+ parameterization="eps", # all assuming fixed variance schedules
71
+ scheduler_config=None,
72
+ use_positional_encodings=False,
73
+ learn_logvar=False,
74
+ logvar_init=0.,
75
+ make_it_fit=False,
76
+ ucg_training=None,
77
+ reset_ema=False,
78
+ reset_num_ema_updates=False,
79
+ ):
80
+ super().__init__()
81
+ assert parameterization in ["eps", "x0", "v"], 'currently only supporting "eps" and "x0" and "v"'
82
+ self.parameterization = parameterization
83
+ print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
84
+ self.cond_stage_model = None
85
+ self.clip_denoised = clip_denoised
86
+ self.log_every_t = log_every_t
87
+ self.first_stage_key = first_stage_key
88
+ self.image_size = image_size # try conv?
89
+ self.channels = channels
90
+ self.use_positional_encodings = use_positional_encodings
91
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
92
+ count_params(self.model, verbose=True)
93
+ self.use_ema = use_ema
94
+ if self.use_ema:
95
+ self.model_ema = LitEma(self.model)
96
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
97
+
98
+ self.use_scheduler = scheduler_config is not None
99
+ if self.use_scheduler:
100
+ self.scheduler_config = scheduler_config
101
+
102
+ self.v_posterior = v_posterior
103
+ self.original_elbo_weight = original_elbo_weight
104
+ self.l_simple_weight = l_simple_weight
105
+
106
+ if monitor is not None:
107
+ self.monitor = monitor
108
+ self.make_it_fit = make_it_fit
109
+ if reset_ema: assert exists(ckpt_path)
110
+ if ckpt_path is not None:
111
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
112
+ if reset_ema:
113
+ assert self.use_ema
114
+ print(f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
115
+ self.model_ema = LitEma(self.model)
116
+ if reset_num_ema_updates:
117
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
118
+ assert self.use_ema
119
+ self.model_ema.reset_num_updates()
120
+
121
+ self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
122
+ linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
123
+
124
+ self.loss_type = loss_type
125
+
126
+ self.learn_logvar = learn_logvar
127
+ self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
128
+ if self.learn_logvar:
129
+ self.logvar = nn.Parameter(self.logvar, requires_grad=True)
130
+
131
+ self.ucg_training = ucg_training or dict()
132
+ if self.ucg_training:
133
+ self.ucg_prng = np.random.RandomState()
134
+
135
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
136
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
137
+ if exists(given_betas):
138
+ betas = given_betas
139
+ else:
140
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
141
+ cosine_s=cosine_s)
142
+ alphas = 1. - betas
143
+ alphas_cumprod = np.cumprod(alphas, axis=0)
144
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
145
+
146
+ timesteps, = betas.shape
147
+ self.num_timesteps = int(timesteps)
148
+ self.linear_start = linear_start
149
+ self.linear_end = linear_end
150
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
151
+
152
+ to_torch = partial(torch.tensor, dtype=torch.float32)
153
+
154
+ self.register_buffer('betas', to_torch(betas))
155
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
156
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
157
+
158
+ # calculations for diffusion q(x_t | x_{t-1}) and others
159
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
160
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
161
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
162
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
163
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
164
+
165
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
166
+ posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
167
+ 1. - alphas_cumprod) + self.v_posterior * betas
168
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
169
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
170
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
171
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
172
+ self.register_buffer('posterior_mean_coef1', to_torch(
173
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
174
+ self.register_buffer('posterior_mean_coef2', to_torch(
175
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
176
+
177
+ if self.parameterization == "eps":
178
+ lvlb_weights = self.betas ** 2 / (
179
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
180
+ elif self.parameterization == "x0":
181
+ lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
182
+ elif self.parameterization == "v":
183
+ lvlb_weights = torch.ones_like(self.betas ** 2 / (
184
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)))
185
+ else:
186
+ raise NotImplementedError("mu not supported")
187
+ lvlb_weights[0] = lvlb_weights[1]
188
+ self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
189
+ assert not torch.isnan(self.lvlb_weights).all()
190
+
191
+ @contextmanager
192
+ def ema_scope(self, context=None):
193
+ if self.use_ema:
194
+ self.model_ema.store(self.model.parameters())
195
+ self.model_ema.copy_to(self.model)
196
+ if context is not None:
197
+ print(f"{context}: Switched to EMA weights")
198
+ try:
199
+ yield None
200
+ finally:
201
+ if self.use_ema:
202
+ self.model_ema.restore(self.model.parameters())
203
+ if context is not None:
204
+ print(f"{context}: Restored training weights")
205
+
206
+ @torch.no_grad()
207
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
208
+ sd = torch.load(path, map_location="cpu")
209
+ if "state_dict" in list(sd.keys()):
210
+ sd = sd["state_dict"]
211
+ keys = list(sd.keys())
212
+ for k in keys:
213
+ for ik in ignore_keys:
214
+ if k.startswith(ik):
215
+ print("Deleting key {} from state_dict.".format(k))
216
+ del sd[k]
217
+ if self.make_it_fit:
218
+ n_params = len([name for name, _ in
219
+ itertools.chain(self.named_parameters(),
220
+ self.named_buffers())])
221
+ for name, param in tqdm(
222
+ itertools.chain(self.named_parameters(),
223
+ self.named_buffers()),
224
+ desc="Fitting old weights to new weights",
225
+ total=n_params
226
+ ):
227
+ if not name in sd:
228
+ continue
229
+ old_shape = sd[name].shape
230
+ new_shape = param.shape
231
+ assert len(old_shape) == len(new_shape)
232
+ if len(new_shape) > 2:
233
+ # we only modify first two axes
234
+ assert new_shape[2:] == old_shape[2:]
235
+ # assumes first axis corresponds to output dim
236
+ if not new_shape == old_shape:
237
+ new_param = param.clone()
238
+ old_param = sd[name]
239
+ if len(new_shape) == 1:
240
+ for i in range(new_param.shape[0]):
241
+ new_param[i] = old_param[i % old_shape[0]]
242
+ elif len(new_shape) >= 2:
243
+ for i in range(new_param.shape[0]):
244
+ for j in range(new_param.shape[1]):
245
+ new_param[i, j] = old_param[i % old_shape[0], j % old_shape[1]]
246
+
247
+ n_used_old = torch.ones(old_shape[1])
248
+ for j in range(new_param.shape[1]):
249
+ n_used_old[j % old_shape[1]] += 1
250
+ n_used_new = torch.zeros(new_shape[1])
251
+ for j in range(new_param.shape[1]):
252
+ n_used_new[j] = n_used_old[j % old_shape[1]]
253
+
254
+ n_used_new = n_used_new[None, :]
255
+ while len(n_used_new.shape) < len(new_shape):
256
+ n_used_new = n_used_new.unsqueeze(-1)
257
+ new_param /= n_used_new
258
+
259
+ sd[name] = new_param
260
+
261
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
262
+ sd, strict=False)
263
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
264
+ if len(missing) > 0:
265
+ print(f"Missing Keys:\n {missing}")
266
+ if len(unexpected) > 0:
267
+ print(f"\nUnexpected Keys:\n {unexpected}")
268
+
269
+ def q_mean_variance(self, x_start, t):
270
+ """
271
+ Get the distribution q(x_t | x_0).
272
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
273
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
274
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
275
+ """
276
+ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
277
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
278
+ log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
279
+ return mean, variance, log_variance
280
+
281
+ def predict_start_from_noise(self, x_t, t, noise):
282
+ return (
283
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
284
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
285
+ )
286
+
287
+ def predict_start_from_z_and_v(self, x_t, t, v):
288
+ # self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
289
+ # self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
290
+ return (
291
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
292
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
293
+ )
294
+
295
+ def predict_eps_from_z_and_v(self, x_t, t, v):
296
+ return (
297
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * v +
298
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * x_t
299
+ )
300
+
301
+ def q_posterior(self, x_start, x_t, t):
302
+ posterior_mean = (
303
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
304
+ extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
305
+ )
306
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
307
+ posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
308
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
309
+
310
+ def p_mean_variance(self, x, t, clip_denoised: bool):
311
+ model_out = self.model(x, t)
312
+ if self.parameterization == "eps":
313
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
314
+ elif self.parameterization == "x0":
315
+ x_recon = model_out
316
+ if clip_denoised:
317
+ x_recon.clamp_(-1., 1.)
318
+
319
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
320
+ return model_mean, posterior_variance, posterior_log_variance
321
+
322
+ @torch.no_grad()
323
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
324
+ b, *_, device = *x.shape, x.device
325
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
326
+ noise = noise_like(x.shape, device, repeat_noise)
327
+ # no noise when t == 0
328
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
329
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
330
+
331
+ @torch.no_grad()
332
+ def p_sample_loop(self, shape, return_intermediates=False):
333
+ device = self.betas.device
334
+ b = shape[0]
335
+ img = torch.randn(shape, device=device)
336
+ intermediates = [img]
337
+ for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
338
+ img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
339
+ clip_denoised=self.clip_denoised)
340
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
341
+ intermediates.append(img)
342
+ if return_intermediates:
343
+ return img, intermediates
344
+ return img
345
+
346
+ @torch.no_grad()
347
+ def sample(self, batch_size=16, return_intermediates=False):
348
+ image_size = self.image_size
349
+ channels = self.channels
350
+ return self.p_sample_loop((batch_size, channels, image_size, image_size),
351
+ return_intermediates=return_intermediates)
352
+
353
+ def q_sample(self, x_start, t, noise=None):
354
+ noise = default(noise, lambda: torch.randn_like(x_start))
355
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
356
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
357
+
358
+ def get_v(self, x, noise, t):
359
+ return (
360
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
361
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
362
+ )
363
+
364
+ def get_loss(self, pred, target, mean=True):
365
+ if self.loss_type == 'l1':
366
+ loss = (target - pred).abs()
367
+ if mean:
368
+ loss = loss.mean()
369
+ elif self.loss_type == 'l2':
370
+ if mean:
371
+ loss = torch.nn.functional.mse_loss(target, pred)
372
+ else:
373
+ loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
374
+ else:
375
+ raise NotImplementedError("unknown loss type '{loss_type}'")
376
+
377
+ return loss
378
+
379
+ def p_losses(self, x_start, t, noise=None):
380
+ noise = default(noise, lambda: torch.randn_like(x_start))
381
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
382
+ model_out = self.model(x_noisy, t)
383
+
384
+ loss_dict = {}
385
+ if self.parameterization == "eps":
386
+ target = noise
387
+ elif self.parameterization == "x0":
388
+ target = x_start
389
+ elif self.parameterization == "v":
390
+ target = self.get_v(x_start, noise, t)
391
+ else:
392
+ raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported")
393
+
394
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
395
+
396
+ log_prefix = 'train' if self.training else 'val'
397
+
398
+ loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
399
+ loss_simple = loss.mean() * self.l_simple_weight
400
+
401
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
402
+ loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
403
+
404
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
405
+
406
+ loss_dict.update({f'{log_prefix}/loss': loss})
407
+
408
+ return loss, loss_dict
409
+
410
+ def forward(self, x, *args, **kwargs):
411
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
412
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
413
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
414
+ return self.p_losses(x, t, *args, **kwargs)
415
+
416
+ def get_input(self, batch, k):
417
+ x = batch[k]
418
+ if len(x.shape) == 3:
419
+ x = x[..., None]
420
+ x = rearrange(x, 'b h w c -> b c h w')
421
+ x = x.to(memory_format=torch.contiguous_format).float()
422
+ return x
423
+
424
+ def shared_step(self, batch):
425
+ x = self.get_input(batch, self.first_stage_key)
426
+ loss, loss_dict = self(x)
427
+ return loss, loss_dict
428
+
429
+ def training_step(self, batch, batch_idx):
430
+ for k in self.ucg_training:
431
+ p = self.ucg_training[k]["p"]
432
+ val = self.ucg_training[k]["val"]
433
+ if val is None:
434
+ val = ""
435
+ for i in range(len(batch[k])):
436
+ if self.ucg_prng.choice(2, p=[1 - p, p]):
437
+ batch[k][i] = val
438
+
439
+ loss, loss_dict = self.shared_step(batch)
440
+
441
+ self.log_dict(loss_dict, prog_bar=True,
442
+ logger=True, on_step=True, on_epoch=True)
443
+
444
+ self.log("global_step", self.global_step,
445
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
446
+
447
+ if self.use_scheduler:
448
+ lr = self.optimizers().param_groups[0]['lr']
449
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
450
+
451
+ return loss
452
+
453
+ @torch.no_grad()
454
+ def validation_step(self, batch, batch_idx):
455
+ _, loss_dict_no_ema = self.shared_step(batch)
456
+ with self.ema_scope():
457
+ _, loss_dict_ema = self.shared_step(batch)
458
+ loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
459
+ self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
460
+ self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
461
+
462
+ def on_train_batch_end(self, *args, **kwargs):
463
+ if self.use_ema:
464
+ self.model_ema(self.model)
465
+
466
+ def _get_rows_from_list(self, samples):
467
+ n_imgs_per_row = len(samples)
468
+ denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
469
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
470
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
471
+ return denoise_grid
472
+
473
+ @torch.no_grad()
474
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
475
+ log = dict()
476
+ x = self.get_input(batch, self.first_stage_key)
477
+ N = min(x.shape[0], N)
478
+ n_row = min(x.shape[0], n_row)
479
+ x = x.to(self.device)[:N]
480
+ log["inputs"] = x
481
+
482
+ # get diffusion row
483
+ diffusion_row = list()
484
+ x_start = x[:n_row]
485
+
486
+ for t in range(self.num_timesteps):
487
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
488
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
489
+ t = t.to(self.device).long()
490
+ noise = torch.randn_like(x_start)
491
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
492
+ diffusion_row.append(x_noisy)
493
+
494
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
495
+
496
+ if sample:
497
+ # get denoise row
498
+ with self.ema_scope("Plotting"):
499
+ samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
500
+
501
+ log["samples"] = samples
502
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
503
+
504
+ if return_keys:
505
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
506
+ return log
507
+ else:
508
+ return {key: log[key] for key in return_keys}
509
+ return log
510
+
511
+ def configure_optimizers(self):
512
+ lr = self.learning_rate
513
+ params = list(self.model.parameters())
514
+ if self.learn_logvar:
515
+ params = params + [self.logvar]
516
+ opt = torch.optim.AdamW(params, lr=lr)
517
+ return opt
518
+
519
+
520
+ class LatentDiffusion(DDPM):
521
+ """main class"""
522
+
523
+ def __init__(self,
524
+ first_stage_config={},
525
+ cond_stage_config={},
526
+ num_timesteps_cond=None,
527
+ cond_stage_key="image",
528
+ cond_stage_trainable=False,
529
+ concat_mode=True,
530
+ cond_stage_forward=None,
531
+ conditioning_key=None,
532
+ scale_factor=1.0,
533
+ scale_by_std=False,
534
+ force_null_conditioning=False,
535
+ *args, **kwargs):
536
+ self.force_null_conditioning = force_null_conditioning
537
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
538
+ self.scale_by_std = scale_by_std
539
+ assert self.num_timesteps_cond <= kwargs['timesteps']
540
+ # for backwards compatibility after implementation of DiffusionWrapper
541
+ if conditioning_key is None:
542
+ conditioning_key = 'concat' if concat_mode else 'crossattn'
543
+ if cond_stage_config == '__is_unconditional__' and not self.force_null_conditioning:
544
+ conditioning_key = None
545
+ ckpt_path = kwargs.pop("ckpt_path", None)
546
+ reset_ema = kwargs.pop("reset_ema", False)
547
+ reset_num_ema_updates = kwargs.pop("reset_num_ema_updates", False)
548
+ ignore_keys = kwargs.pop("ignore_keys", [])
549
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
550
+ self.concat_mode = concat_mode
551
+ self.cond_stage_trainable = cond_stage_trainable
552
+ self.cond_stage_key = cond_stage_key
553
+ try:
554
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
555
+ except:
556
+ self.num_downs = 0
557
+ if not scale_by_std:
558
+ self.scale_factor = scale_factor
559
+ else:
560
+ self.register_buffer('scale_factor', torch.tensor(scale_factor))
561
+
562
+ # self.instantiate_first_stage(first_stage_config)
563
+ # self.instantiate_cond_stage(cond_stage_config)
564
+
565
+ self.cond_stage_forward = cond_stage_forward
566
+ self.clip_denoised = False
567
+ self.bbox_tokenizer = None
568
+
569
+ self.restarted_from_ckpt = False
570
+ if ckpt_path is not None:
571
+ self.init_from_ckpt(ckpt_path, ignore_keys)
572
+ self.restarted_from_ckpt = True
573
+ if reset_ema:
574
+ assert self.use_ema
575
+ print(
576
+ f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.")
577
+ self.model_ema = LitEma(self.model)
578
+ if reset_num_ema_updates:
579
+ print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ")
580
+ assert self.use_ema
581
+ self.model_ema.reset_num_updates()
582
+
583
+ def make_cond_schedule(self, ):
584
+ self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
585
+ ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
586
+ self.cond_ids[:self.num_timesteps_cond] = ids
587
+
588
+ # @rank_zero_only
589
+ @torch.no_grad()
590
+ def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
591
+ # only for very first batch
592
+ if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
593
+ assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
594
+ # set rescale weight to 1./std of encodings
595
+ print("### USING STD-RESCALING ###")
596
+ x = super().get_input(batch, self.first_stage_key)
597
+ x = x.to(self.device)
598
+ encoder_posterior = self.encode_first_stage(x)
599
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
600
+ del self.scale_factor
601
+ self.register_buffer('scale_factor', 1. / z.flatten().std())
602
+ print(f"setting self.scale_factor to {self.scale_factor}")
603
+ print("### USING STD-RESCALING ###")
604
+
605
+ def register_schedule(self,
606
+ given_betas=None, beta_schedule="linear", timesteps=1000,
607
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
608
+ super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
609
+
610
+ self.shorten_cond_schedule = self.num_timesteps_cond > 1
611
+ if self.shorten_cond_schedule:
612
+ self.make_cond_schedule()
613
+
614
+ def instantiate_first_stage(self, config):
615
+ model = instantiate_from_config(config)
616
+ self.first_stage_model = model.eval()
617
+ self.first_stage_model.train = disabled_train
618
+ for param in self.first_stage_model.parameters():
619
+ param.requires_grad = False
620
+
621
+ def instantiate_cond_stage(self, config):
622
+ if not self.cond_stage_trainable:
623
+ if config == "__is_first_stage__":
624
+ print("Using first stage also as cond stage.")
625
+ self.cond_stage_model = self.first_stage_model
626
+ elif config == "__is_unconditional__":
627
+ print(f"Training {self.__class__.__name__} as an unconditional model.")
628
+ self.cond_stage_model = None
629
+ # self.be_unconditional = True
630
+ else:
631
+ model = instantiate_from_config(config)
632
+ self.cond_stage_model = model.eval()
633
+ self.cond_stage_model.train = disabled_train
634
+ for param in self.cond_stage_model.parameters():
635
+ param.requires_grad = False
636
+ else:
637
+ assert config != '__is_first_stage__'
638
+ assert config != '__is_unconditional__'
639
+ model = instantiate_from_config(config)
640
+ self.cond_stage_model = model
641
+
642
+ def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
643
+ denoise_row = []
644
+ for zd in tqdm(samples, desc=desc):
645
+ denoise_row.append(self.decode_first_stage(zd.to(self.device),
646
+ force_not_quantize=force_no_decoder_quantization))
647
+ n_imgs_per_row = len(denoise_row)
648
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
649
+ denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
650
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
651
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
652
+ return denoise_grid
653
+
654
+ def get_first_stage_encoding(self, encoder_posterior):
655
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
656
+ z = encoder_posterior.sample()
657
+ elif isinstance(encoder_posterior, torch.Tensor):
658
+ z = encoder_posterior
659
+ else:
660
+ raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
661
+ return self.scale_factor * z
662
+
663
+ def get_learned_conditioning(self, c):
664
+ if self.cond_stage_forward is None:
665
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
666
+ c = self.cond_stage_model.encode(c)
667
+ if isinstance(c, DiagonalGaussianDistribution):
668
+ c = c.mode()
669
+ else:
670
+ c = self.cond_stage_model(c)
671
+ else:
672
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
673
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
674
+ return c
675
+
676
+ def meshgrid(self, h, w):
677
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
678
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
679
+
680
+ arr = torch.cat([y, x], dim=-1)
681
+ return arr
682
+
683
+ def delta_border(self, h, w):
684
+ """
685
+ :param h: height
686
+ :param w: width
687
+ :return: normalized distance to image border,
688
+ wtith min distance = 0 at border and max dist = 0.5 at image center
689
+ """
690
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
691
+ arr = self.meshgrid(h, w) / lower_right_corner
692
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
693
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
694
+ edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
695
+ return edge_dist
696
+
697
+ def get_weighting(self, h, w, Ly, Lx, device):
698
+ weighting = self.delta_border(h, w)
699
+ weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
700
+ self.split_input_params["clip_max_weight"], )
701
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
702
+
703
+ if self.split_input_params["tie_braker"]:
704
+ L_weighting = self.delta_border(Ly, Lx)
705
+ L_weighting = torch.clip(L_weighting,
706
+ self.split_input_params["clip_min_tie_weight"],
707
+ self.split_input_params["clip_max_tie_weight"])
708
+
709
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
710
+ weighting = weighting * L_weighting
711
+ return weighting
712
+
713
+ def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code
714
+ """
715
+ :param x: img of size (bs, c, h, w)
716
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
717
+ """
718
+ bs, nc, h, w = x.shape
719
+
720
+ # number of crops in image
721
+ Ly = (h - kernel_size[0]) // stride[0] + 1
722
+ Lx = (w - kernel_size[1]) // stride[1] + 1
723
+
724
+ if uf == 1 and df == 1:
725
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
726
+ unfold = torch.nn.Unfold(**fold_params)
727
+
728
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
729
+
730
+ weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
731
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
732
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
733
+
734
+ elif uf > 1 and df == 1:
735
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
736
+ unfold = torch.nn.Unfold(**fold_params)
737
+
738
+ fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
739
+ dilation=1, padding=0,
740
+ stride=(stride[0] * uf, stride[1] * uf))
741
+ fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
742
+
743
+ weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
744
+ normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap
745
+ weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
746
+
747
+ elif df > 1 and uf == 1:
748
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
749
+ unfold = torch.nn.Unfold(**fold_params)
750
+
751
+ fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
752
+ dilation=1, padding=0,
753
+ stride=(stride[0] // df, stride[1] // df))
754
+ fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
755
+
756
+ weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
757
+ normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap
758
+ weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
759
+
760
+ else:
761
+ raise NotImplementedError
762
+
763
+ return fold, unfold, normalization, weighting
764
+
765
+ @torch.no_grad()
766
+ def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
767
+ cond_key=None, return_original_cond=False, bs=None, return_x=False):
768
+ x = super().get_input(batch, k)
769
+ if bs is not None:
770
+ x = x[:bs]
771
+ x = x.to(self.device)
772
+ encoder_posterior = self.encode_first_stage(x)
773
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
774
+
775
+ if self.model.conditioning_key is not None and not self.force_null_conditioning:
776
+ if cond_key is None:
777
+ cond_key = self.cond_stage_key
778
+ if cond_key != self.first_stage_key:
779
+ if cond_key in ['caption', 'coordinates_bbox', "txt"]:
780
+ xc = batch[cond_key]
781
+ elif cond_key in ['class_label', 'cls']:
782
+ xc = batch
783
+ else:
784
+ xc = super().get_input(batch, cond_key).to(self.device)
785
+ else:
786
+ xc = x
787
+ if not self.cond_stage_trainable or force_c_encode:
788
+ if isinstance(xc, dict) or isinstance(xc, list):
789
+ c = self.get_learned_conditioning(xc)
790
+ else:
791
+ c = self.get_learned_conditioning(xc.to(self.device))
792
+ else:
793
+ c = xc
794
+ if bs is not None:
795
+ c = c[:bs]
796
+
797
+ if self.use_positional_encodings:
798
+ pos_x, pos_y = self.compute_latent_shifts(batch)
799
+ ckey = __conditioning_keys__[self.model.conditioning_key]
800
+ c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}
801
+
802
+ else:
803
+ c = None
804
+ xc = None
805
+ if self.use_positional_encodings:
806
+ pos_x, pos_y = self.compute_latent_shifts(batch)
807
+ c = {'pos_x': pos_x, 'pos_y': pos_y}
808
+ out = [z, c]
809
+ if return_first_stage_outputs:
810
+ xrec = self.decode_first_stage(z)
811
+ out.extend([x, xrec])
812
+ if return_x:
813
+ out.extend([x])
814
+ if return_original_cond:
815
+ out.append(xc)
816
+ return out
817
+
818
+ @torch.no_grad()
819
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
820
+ if predict_cids:
821
+ if z.dim() == 4:
822
+ z = torch.argmax(z.exp(), dim=1).long()
823
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
824
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
825
+
826
+ z = 1. / self.scale_factor * z
827
+ return self.first_stage_model.decode(z)
828
+
829
+ @torch.no_grad()
830
+ def encode_first_stage(self, x):
831
+ return self.first_stage_model.encode(x)
832
+
833
+ def shared_step(self, batch, **kwargs):
834
+ x, c = self.get_input(batch, self.first_stage_key)
835
+ loss = self(x, c)
836
+ return loss
837
+
838
+ def forward(self, x, c, *args, **kwargs):
839
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
840
+ if self.model.conditioning_key is not None:
841
+ assert c is not None
842
+ if self.cond_stage_trainable:
843
+ c = self.get_learned_conditioning(c)
844
+ if self.shorten_cond_schedule: # TODO: drop this option
845
+ tc = self.cond_ids[t].to(self.device)
846
+ c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
847
+ return self.p_losses(x, c, t, *args, **kwargs)
848
+
849
+ def apply_model(self, x_noisy, t, cond, return_ids=False):
850
+ if isinstance(cond, dict):
851
+ # hybrid case, cond is expected to be a dict
852
+ pass
853
+ else:
854
+ if not isinstance(cond, list):
855
+ cond = [cond]
856
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
857
+ cond = {key: cond}
858
+
859
+ x_recon = self.model(x_noisy, t, **cond)
860
+
861
+ if isinstance(x_recon, tuple) and not return_ids:
862
+ return x_recon[0]
863
+ else:
864
+ return x_recon
865
+
866
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
867
+ return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
868
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
869
+
870
+ def _prior_bpd(self, x_start):
871
+ """
872
+ Get the prior KL term for the variational lower-bound, measured in
873
+ bits-per-dim.
874
+ This term can't be optimized, as it only depends on the encoder.
875
+ :param x_start: the [N x C x ...] tensor of inputs.
876
+ :return: a batch of [N] KL values (in bits), one per batch element.
877
+ """
878
+ batch_size = x_start.shape[0]
879
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
880
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
881
+ kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
882
+ return mean_flat(kl_prior) / np.log(2.0)
883
+
884
+ def p_losses(self, x_start, cond, t, noise=None):
885
+ noise = default(noise, lambda: torch.randn_like(x_start))
886
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
887
+ model_output = self.apply_model(x_noisy, t, cond)
888
+
889
+ loss_dict = {}
890
+ prefix = 'train' if self.training else 'val'
891
+
892
+ if self.parameterization == "x0":
893
+ target = x_start
894
+ elif self.parameterization == "eps":
895
+ target = noise
896
+ elif self.parameterization == "v":
897
+ target = self.get_v(x_start, noise, t)
898
+ else:
899
+ raise NotImplementedError()
900
+
901
+ loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
902
+ loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
903
+
904
+ logvar_t = self.logvar[t].to(self.device)
905
+ loss = loss_simple / torch.exp(logvar_t) + logvar_t
906
+ # loss = loss_simple / torch.exp(self.logvar) + self.logvar
907
+ if self.learn_logvar:
908
+ loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
909
+ loss_dict.update({'logvar': self.logvar.data.mean()})
910
+
911
+ loss = self.l_simple_weight * loss.mean()
912
+
913
+ loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
914
+ loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
915
+ loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
916
+ loss += (self.original_elbo_weight * loss_vlb)
917
+ loss_dict.update({f'{prefix}/loss': loss})
918
+
919
+ return loss, loss_dict
920
+
921
+ def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
922
+ return_x0=False, score_corrector=None, corrector_kwargs=None):
923
+ t_in = t
924
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
925
+
926
+ if score_corrector is not None:
927
+ assert self.parameterization == "eps"
928
+ model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
929
+
930
+ if return_codebook_ids:
931
+ model_out, logits = model_out
932
+
933
+ if self.parameterization == "eps":
934
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
935
+ elif self.parameterization == "x0":
936
+ x_recon = model_out
937
+ else:
938
+ raise NotImplementedError()
939
+
940
+ if clip_denoised:
941
+ x_recon.clamp_(-1., 1.)
942
+ if quantize_denoised:
943
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
944
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
945
+ if return_codebook_ids:
946
+ return model_mean, posterior_variance, posterior_log_variance, logits
947
+ elif return_x0:
948
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
949
+ else:
950
+ return model_mean, posterior_variance, posterior_log_variance
951
+
952
+ @torch.no_grad()
953
+ def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
954
+ return_codebook_ids=False, quantize_denoised=False, return_x0=False,
955
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
956
+ b, *_, device = *x.shape, x.device
957
+ outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
958
+ return_codebook_ids=return_codebook_ids,
959
+ quantize_denoised=quantize_denoised,
960
+ return_x0=return_x0,
961
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
962
+ if return_codebook_ids:
963
+ raise DeprecationWarning("Support dropped.")
964
+ model_mean, _, model_log_variance, logits = outputs
965
+ elif return_x0:
966
+ model_mean, _, model_log_variance, x0 = outputs
967
+ else:
968
+ model_mean, _, model_log_variance = outputs
969
+
970
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
971
+ if noise_dropout > 0.:
972
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
973
+ # no noise when t == 0
974
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
975
+
976
+ if return_codebook_ids:
977
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
978
+ if return_x0:
979
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
980
+ else:
981
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
982
+
983
+ @torch.no_grad()
984
+ def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
985
+ img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
986
+ score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
987
+ log_every_t=None):
988
+ if not log_every_t:
989
+ log_every_t = self.log_every_t
990
+ timesteps = self.num_timesteps
991
+ if batch_size is not None:
992
+ b = batch_size if batch_size is not None else shape[0]
993
+ shape = [batch_size] + list(shape)
994
+ else:
995
+ b = batch_size = shape[0]
996
+ if x_T is None:
997
+ img = torch.randn(shape, device=self.device)
998
+ else:
999
+ img = x_T
1000
+ intermediates = []
1001
+ if cond is not None:
1002
+ if isinstance(cond, dict):
1003
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1004
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1005
+ else:
1006
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1007
+
1008
+ if start_T is not None:
1009
+ timesteps = min(timesteps, start_T)
1010
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
1011
+ total=timesteps) if verbose else reversed(
1012
+ range(0, timesteps))
1013
+ if type(temperature) == float:
1014
+ temperature = [temperature] * timesteps
1015
+
1016
+ for i in iterator:
1017
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
1018
+ if self.shorten_cond_schedule:
1019
+ assert self.model.conditioning_key != 'hybrid'
1020
+ tc = self.cond_ids[ts].to(cond.device)
1021
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1022
+
1023
+ img, x0_partial = self.p_sample(img, cond, ts,
1024
+ clip_denoised=self.clip_denoised,
1025
+ quantize_denoised=quantize_denoised, return_x0=True,
1026
+ temperature=temperature[i], noise_dropout=noise_dropout,
1027
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
1028
+ if mask is not None:
1029
+ assert x0 is not None
1030
+ img_orig = self.q_sample(x0, ts)
1031
+ img = img_orig * mask + (1. - mask) * img
1032
+
1033
+ if i % log_every_t == 0 or i == timesteps - 1:
1034
+ intermediates.append(x0_partial)
1035
+ if callback: callback(i)
1036
+ if img_callback: img_callback(img, i)
1037
+ return img, intermediates
1038
+
1039
+ @torch.no_grad()
1040
+ def p_sample_loop(self, cond, shape, return_intermediates=False,
1041
+ x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
1042
+ mask=None, x0=None, img_callback=None, start_T=None,
1043
+ log_every_t=None):
1044
+
1045
+ if not log_every_t:
1046
+ log_every_t = self.log_every_t
1047
+ device = self.betas.device
1048
+ b = shape[0]
1049
+ if x_T is None:
1050
+ img = torch.randn(shape, device=device)
1051
+ else:
1052
+ img = x_T
1053
+
1054
+ intermediates = [img]
1055
+ if timesteps is None:
1056
+ timesteps = self.num_timesteps
1057
+
1058
+ if start_T is not None:
1059
+ timesteps = min(timesteps, start_T)
1060
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
1061
+ range(0, timesteps))
1062
+
1063
+ if mask is not None:
1064
+ assert x0 is not None
1065
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
1066
+
1067
+ for i in iterator:
1068
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
1069
+ if self.shorten_cond_schedule:
1070
+ assert self.model.conditioning_key != 'hybrid'
1071
+ tc = self.cond_ids[ts].to(cond.device)
1072
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
1073
+
1074
+ img = self.p_sample(img, cond, ts,
1075
+ clip_denoised=self.clip_denoised,
1076
+ quantize_denoised=quantize_denoised)
1077
+ if mask is not None:
1078
+ img_orig = self.q_sample(x0, ts)
1079
+ img = img_orig * mask + (1. - mask) * img
1080
+
1081
+ if i % log_every_t == 0 or i == timesteps - 1:
1082
+ intermediates.append(img)
1083
+ if callback: callback(i)
1084
+ if img_callback: img_callback(img, i)
1085
+
1086
+ if return_intermediates:
1087
+ return img, intermediates
1088
+ return img
1089
+
1090
+ @torch.no_grad()
1091
+ def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
1092
+ verbose=True, timesteps=None, quantize_denoised=False,
1093
+ mask=None, x0=None, shape=None, **kwargs):
1094
+ if shape is None:
1095
+ shape = (batch_size, self.channels, self.image_size, self.image_size)
1096
+ if cond is not None:
1097
+ if isinstance(cond, dict):
1098
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
1099
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
1100
+ else:
1101
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
1102
+ return self.p_sample_loop(cond,
1103
+ shape,
1104
+ return_intermediates=return_intermediates, x_T=x_T,
1105
+ verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
1106
+ mask=mask, x0=x0)
1107
+
1108
+ @torch.no_grad()
1109
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
1110
+ if ddim:
1111
+ ddim_sampler = DDIMSampler(self)
1112
+ shape = (self.channels, self.image_size, self.image_size)
1113
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size,
1114
+ shape, cond, verbose=False, **kwargs)
1115
+
1116
+ else:
1117
+ samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
1118
+ return_intermediates=True, **kwargs)
1119
+
1120
+ return samples, intermediates
1121
+
1122
+ @torch.no_grad()
1123
+ def get_unconditional_conditioning(self, batch_size, null_label=None):
1124
+ if null_label is not None:
1125
+ xc = null_label
1126
+ # if isinstance(xc, ListConfig):
1127
+ # xc = list(xc)
1128
+ if isinstance(xc, dict) or isinstance(xc, list):
1129
+ c = self.get_learned_conditioning(xc)
1130
+ else:
1131
+ if hasattr(xc, "to"):
1132
+ xc = xc.to(self.device)
1133
+ c = self.get_learned_conditioning(xc)
1134
+ else:
1135
+ if self.cond_stage_key in ["class_label", "cls"]:
1136
+ xc = self.cond_stage_model.get_unconditional_conditioning(batch_size, device=self.device)
1137
+ return self.get_learned_conditioning(xc)
1138
+ else:
1139
+ raise NotImplementedError("todo")
1140
+ if isinstance(c, list): # in case the encoder gives us a list
1141
+ for i in range(len(c)):
1142
+ c[i] = repeat(c[i], '1 ... -> b ...', b=batch_size).to(self.device)
1143
+ else:
1144
+ c = repeat(c, '1 ... -> b ...', b=batch_size).to(self.device)
1145
+ return c
1146
+
1147
+ @torch.no_grad()
1148
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=50, ddim_eta=0., return_keys=None,
1149
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1150
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1151
+ use_ema_scope=True,
1152
+ **kwargs):
1153
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1154
+ use_ddim = ddim_steps is not None
1155
+
1156
+ log = dict()
1157
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
1158
+ return_first_stage_outputs=True,
1159
+ force_c_encode=True,
1160
+ return_original_cond=True,
1161
+ bs=N)
1162
+ N = min(x.shape[0], N)
1163
+ n_row = min(x.shape[0], n_row)
1164
+ log["inputs"] = x
1165
+ log["reconstruction"] = xrec
1166
+ if self.model.conditioning_key is not None:
1167
+ if hasattr(self.cond_stage_model, "decode"):
1168
+ xc = self.cond_stage_model.decode(c)
1169
+ log["conditioning"] = xc
1170
+ elif self.cond_stage_key in ["caption", "txt"]:
1171
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1172
+ log["conditioning"] = xc
1173
+ elif self.cond_stage_key in ['class_label', "cls"]:
1174
+ try:
1175
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1176
+ log['conditioning'] = xc
1177
+ except KeyError:
1178
+ # probably no "human_label" in batch
1179
+ pass
1180
+ elif isimage(xc):
1181
+ log["conditioning"] = xc
1182
+ if ismap(xc):
1183
+ log["original_conditioning"] = self.to_rgb(xc)
1184
+
1185
+ if plot_diffusion_rows:
1186
+ # get diffusion row
1187
+ diffusion_row = list()
1188
+ z_start = z[:n_row]
1189
+ for t in range(self.num_timesteps):
1190
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1191
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1192
+ t = t.to(self.device).long()
1193
+ noise = torch.randn_like(z_start)
1194
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1195
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1196
+
1197
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1198
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1199
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1200
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1201
+ log["diffusion_row"] = diffusion_grid
1202
+
1203
+ if sample:
1204
+ # get denoise row
1205
+ with ema_scope("Sampling"):
1206
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1207
+ ddim_steps=ddim_steps, eta=ddim_eta)
1208
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1209
+ x_samples = self.decode_first_stage(samples)
1210
+ log["samples"] = x_samples
1211
+ if plot_denoise_rows:
1212
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1213
+ log["denoise_row"] = denoise_grid
1214
+
1215
+ if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
1216
+ self.first_stage_model, IdentityFirstStage):
1217
+ # also display when quantizing x0 while sampling
1218
+ with ema_scope("Plotting Quantized Denoised"):
1219
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1220
+ ddim_steps=ddim_steps, eta=ddim_eta,
1221
+ quantize_denoised=True)
1222
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
1223
+ # quantize_denoised=True)
1224
+ x_samples = self.decode_first_stage(samples.to(self.device))
1225
+ log["samples_x0_quantized"] = x_samples
1226
+
1227
+ if unconditional_guidance_scale > 1.0:
1228
+ uc = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1229
+ if self.model.conditioning_key == "crossattn-adm":
1230
+ uc = {"c_crossattn": [uc], "c_adm": c["c_adm"]}
1231
+ with ema_scope("Sampling with classifier-free guidance"):
1232
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1233
+ ddim_steps=ddim_steps, eta=ddim_eta,
1234
+ unconditional_guidance_scale=unconditional_guidance_scale,
1235
+ unconditional_conditioning=uc,
1236
+ )
1237
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1238
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1239
+
1240
+ if inpaint:
1241
+ # make a simple center square
1242
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
1243
+ mask = torch.ones(N, h, w).to(self.device)
1244
+ # zeros will be filled in
1245
+ mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
1246
+ mask = mask[:, None, ...]
1247
+ with ema_scope("Plotting Inpaint"):
1248
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1249
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1250
+ x_samples = self.decode_first_stage(samples.to(self.device))
1251
+ log["samples_inpainting"] = x_samples
1252
+ log["mask"] = mask
1253
+
1254
+ # outpaint
1255
+ mask = 1. - mask
1256
+ with ema_scope("Plotting Outpaint"):
1257
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
1258
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
1259
+ x_samples = self.decode_first_stage(samples.to(self.device))
1260
+ log["samples_outpainting"] = x_samples
1261
+
1262
+ if plot_progressive_rows:
1263
+ with ema_scope("Plotting Progressives"):
1264
+ img, progressives = self.progressive_denoising(c,
1265
+ shape=(self.channels, self.image_size, self.image_size),
1266
+ batch_size=N)
1267
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1268
+ log["progressive_row"] = prog_row
1269
+
1270
+ if return_keys:
1271
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
1272
+ return log
1273
+ else:
1274
+ return {key: log[key] for key in return_keys}
1275
+ return log
1276
+
1277
+ def configure_optimizers(self):
1278
+ lr = self.learning_rate
1279
+ params = list(self.model.parameters())
1280
+ if self.cond_stage_trainable:
1281
+ print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
1282
+ params = params + list(self.cond_stage_model.parameters())
1283
+ if self.learn_logvar:
1284
+ print('Diffusion model optimizing logvar')
1285
+ params.append(self.logvar)
1286
+ opt = torch.optim.AdamW(params, lr=lr)
1287
+ if self.use_scheduler:
1288
+ assert 'target' in self.scheduler_config
1289
+ scheduler = instantiate_from_config(self.scheduler_config)
1290
+
1291
+ print("Setting up LambdaLR scheduler...")
1292
+ scheduler = [
1293
+ {
1294
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
1295
+ 'interval': 'step',
1296
+ 'frequency': 1
1297
+ }]
1298
+ return [opt], scheduler
1299
+ return opt
1300
+
1301
+ @torch.no_grad()
1302
+ def to_rgb(self, x):
1303
+ x = x.float()
1304
+ if not hasattr(self, "colorize"):
1305
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
1306
+ x = nn.functional.conv2d(x, weight=self.colorize)
1307
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
1308
+ return x
1309
+
1310
+
1311
+ # class DiffusionWrapper(pl.LightningModule):
1312
+ class DiffusionWrapper(torch.nn.Module):
1313
+ def __init__(self, diff_model_config, conditioning_key):
1314
+ super().__init__()
1315
+ self.sequential_cross_attn = diff_model_config.pop("sequential_crossattn", False)
1316
+ self.diffusion_model = instantiate_from_config(diff_model_config)
1317
+ self.conditioning_key = conditioning_key
1318
+ assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm']
1319
+
1320
+ def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None, control=None, transformer_options={}):
1321
+ if self.conditioning_key is None:
1322
+ out = self.diffusion_model(x, t, control=control, transformer_options=transformer_options)
1323
+ elif self.conditioning_key == 'concat':
1324
+ xc = torch.cat([x] + c_concat, dim=1)
1325
+ out = self.diffusion_model(xc, t, control=control, transformer_options=transformer_options)
1326
+ elif self.conditioning_key == 'crossattn':
1327
+ if not self.sequential_cross_attn:
1328
+ cc = torch.cat(c_crossattn, 1)
1329
+ else:
1330
+ cc = c_crossattn
1331
+ if hasattr(self, "scripted_diffusion_model"):
1332
+ # TorchScript changes names of the arguments
1333
+ # with argument cc defined as context=cc scripted model will produce
1334
+ # an error: RuntimeError: forward() is missing value for argument 'argument_3'.
1335
+ out = self.scripted_diffusion_model(x, t, cc, control=control, transformer_options=transformer_options)
1336
+ else:
1337
+ out = self.diffusion_model(x, t, context=cc, control=control, transformer_options=transformer_options)
1338
+ elif self.conditioning_key == 'hybrid':
1339
+ xc = torch.cat([x] + c_concat, dim=1)
1340
+ cc = torch.cat(c_crossattn, 1)
1341
+ out = self.diffusion_model(xc, t, context=cc, control=control, transformer_options=transformer_options)
1342
+ elif self.conditioning_key == 'hybrid-adm':
1343
+ assert c_adm is not None
1344
+ xc = torch.cat([x] + c_concat, dim=1)
1345
+ cc = torch.cat(c_crossattn, 1)
1346
+ out = self.diffusion_model(xc, t, context=cc, y=c_adm, control=control, transformer_options=transformer_options)
1347
+ elif self.conditioning_key == 'crossattn-adm':
1348
+ assert c_adm is not None
1349
+ cc = torch.cat(c_crossattn, 1)
1350
+ out = self.diffusion_model(x, t, context=cc, y=c_adm, control=control, transformer_options=transformer_options)
1351
+ elif self.conditioning_key == 'adm':
1352
+ cc = c_crossattn[0]
1353
+ out = self.diffusion_model(x, t, y=cc, control=control, transformer_options=transformer_options)
1354
+ else:
1355
+ raise NotImplementedError()
1356
+
1357
+ return out
1358
+
1359
+
1360
+ class LatentUpscaleDiffusion(LatentDiffusion):
1361
+ def __init__(self, *args, low_scale_config, low_scale_key="LR", noise_level_key=None, **kwargs):
1362
+ super().__init__(*args, **kwargs)
1363
+ # assumes that neither the cond_stage nor the low_scale_model contain trainable params
1364
+ assert not self.cond_stage_trainable
1365
+ self.instantiate_low_stage(low_scale_config)
1366
+ self.low_scale_key = low_scale_key
1367
+ self.noise_level_key = noise_level_key
1368
+
1369
+ def instantiate_low_stage(self, config):
1370
+ model = instantiate_from_config(config)
1371
+ self.low_scale_model = model.eval()
1372
+ self.low_scale_model.train = disabled_train
1373
+ for param in self.low_scale_model.parameters():
1374
+ param.requires_grad = False
1375
+
1376
+ @torch.no_grad()
1377
+ def get_input(self, batch, k, cond_key=None, bs=None, log_mode=False):
1378
+ if not log_mode:
1379
+ z, c = super().get_input(batch, k, force_c_encode=True, bs=bs)
1380
+ else:
1381
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1382
+ force_c_encode=True, return_original_cond=True, bs=bs)
1383
+ x_low = batch[self.low_scale_key][:bs]
1384
+ x_low = rearrange(x_low, 'b h w c -> b c h w')
1385
+ x_low = x_low.to(memory_format=torch.contiguous_format).float()
1386
+ zx, noise_level = self.low_scale_model(x_low)
1387
+ if self.noise_level_key is not None:
1388
+ # get noise level from batch instead, e.g. when extracting a custom noise level for bsr
1389
+ raise NotImplementedError('TODO')
1390
+
1391
+ all_conds = {"c_concat": [zx], "c_crossattn": [c], "c_adm": noise_level}
1392
+ if log_mode:
1393
+ # TODO: maybe disable if too expensive
1394
+ x_low_rec = self.low_scale_model.decode(zx)
1395
+ return z, all_conds, x, xrec, xc, x_low, x_low_rec, noise_level
1396
+ return z, all_conds
1397
+
1398
+ @torch.no_grad()
1399
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
1400
+ plot_denoise_rows=False, plot_progressive_rows=True, plot_diffusion_rows=True,
1401
+ unconditional_guidance_scale=1., unconditional_guidance_label=None, use_ema_scope=True,
1402
+ **kwargs):
1403
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1404
+ use_ddim = ddim_steps is not None
1405
+
1406
+ log = dict()
1407
+ z, c, x, xrec, xc, x_low, x_low_rec, noise_level = self.get_input(batch, self.first_stage_key, bs=N,
1408
+ log_mode=True)
1409
+ N = min(x.shape[0], N)
1410
+ n_row = min(x.shape[0], n_row)
1411
+ log["inputs"] = x
1412
+ log["reconstruction"] = xrec
1413
+ log["x_lr"] = x_low
1414
+ log[f"x_lr_rec_@noise_levels{'-'.join(map(lambda x: str(x), list(noise_level.cpu().numpy())))}"] = x_low_rec
1415
+ if self.model.conditioning_key is not None:
1416
+ if hasattr(self.cond_stage_model, "decode"):
1417
+ xc = self.cond_stage_model.decode(c)
1418
+ log["conditioning"] = xc
1419
+ elif self.cond_stage_key in ["caption", "txt"]:
1420
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1421
+ log["conditioning"] = xc
1422
+ elif self.cond_stage_key in ['class_label', 'cls']:
1423
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1424
+ log['conditioning'] = xc
1425
+ elif isimage(xc):
1426
+ log["conditioning"] = xc
1427
+ if ismap(xc):
1428
+ log["original_conditioning"] = self.to_rgb(xc)
1429
+
1430
+ if plot_diffusion_rows:
1431
+ # get diffusion row
1432
+ diffusion_row = list()
1433
+ z_start = z[:n_row]
1434
+ for t in range(self.num_timesteps):
1435
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1436
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1437
+ t = t.to(self.device).long()
1438
+ noise = torch.randn_like(z_start)
1439
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1440
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1441
+
1442
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1443
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1444
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1445
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1446
+ log["diffusion_row"] = diffusion_grid
1447
+
1448
+ if sample:
1449
+ # get denoise row
1450
+ with ema_scope("Sampling"):
1451
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1452
+ ddim_steps=ddim_steps, eta=ddim_eta)
1453
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1454
+ x_samples = self.decode_first_stage(samples)
1455
+ log["samples"] = x_samples
1456
+ if plot_denoise_rows:
1457
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1458
+ log["denoise_row"] = denoise_grid
1459
+
1460
+ if unconditional_guidance_scale > 1.0:
1461
+ uc_tmp = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1462
+ # TODO explore better "unconditional" choices for the other keys
1463
+ # maybe guide away from empty text label and highest noise level and maximally degraded zx?
1464
+ uc = dict()
1465
+ for k in c:
1466
+ if k == "c_crossattn":
1467
+ assert isinstance(c[k], list) and len(c[k]) == 1
1468
+ uc[k] = [uc_tmp]
1469
+ elif k == "c_adm": # todo: only run with text-based guidance?
1470
+ assert isinstance(c[k], torch.Tensor)
1471
+ #uc[k] = torch.ones_like(c[k]) * self.low_scale_model.max_noise_level
1472
+ uc[k] = c[k]
1473
+ elif isinstance(c[k], list):
1474
+ uc[k] = [c[k][i] for i in range(len(c[k]))]
1475
+ else:
1476
+ uc[k] = c[k]
1477
+
1478
+ with ema_scope("Sampling with classifier-free guidance"):
1479
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
1480
+ ddim_steps=ddim_steps, eta=ddim_eta,
1481
+ unconditional_guidance_scale=unconditional_guidance_scale,
1482
+ unconditional_conditioning=uc,
1483
+ )
1484
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1485
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1486
+
1487
+ if plot_progressive_rows:
1488
+ with ema_scope("Plotting Progressives"):
1489
+ img, progressives = self.progressive_denoising(c,
1490
+ shape=(self.channels, self.image_size, self.image_size),
1491
+ batch_size=N)
1492
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
1493
+ log["progressive_row"] = prog_row
1494
+
1495
+ return log
1496
+
1497
+
1498
+ class LatentFinetuneDiffusion(LatentDiffusion):
1499
+ """
1500
+ Basis for different finetunas, such as inpainting or depth2image
1501
+ To disable finetuning mode, set finetune_keys to None
1502
+ """
1503
+
1504
+ def __init__(self,
1505
+ concat_keys: tuple,
1506
+ finetune_keys=("model.diffusion_model.input_blocks.0.0.weight",
1507
+ "model_ema.diffusion_modelinput_blocks00weight"
1508
+ ),
1509
+ keep_finetune_dims=4,
1510
+ # if model was trained without concat mode before and we would like to keep these channels
1511
+ c_concat_log_start=None, # to log reconstruction of c_concat codes
1512
+ c_concat_log_end=None,
1513
+ *args, **kwargs
1514
+ ):
1515
+ ckpt_path = kwargs.pop("ckpt_path", None)
1516
+ ignore_keys = kwargs.pop("ignore_keys", list())
1517
+ super().__init__(*args, **kwargs)
1518
+ self.finetune_keys = finetune_keys
1519
+ self.concat_keys = concat_keys
1520
+ self.keep_dims = keep_finetune_dims
1521
+ self.c_concat_log_start = c_concat_log_start
1522
+ self.c_concat_log_end = c_concat_log_end
1523
+ if exists(self.finetune_keys): assert exists(ckpt_path), 'can only finetune from a given checkpoint'
1524
+ if exists(ckpt_path):
1525
+ self.init_from_ckpt(ckpt_path, ignore_keys)
1526
+
1527
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
1528
+ sd = torch.load(path, map_location="cpu")
1529
+ if "state_dict" in list(sd.keys()):
1530
+ sd = sd["state_dict"]
1531
+ keys = list(sd.keys())
1532
+ for k in keys:
1533
+ for ik in ignore_keys:
1534
+ if k.startswith(ik):
1535
+ print("Deleting key {} from state_dict.".format(k))
1536
+ del sd[k]
1537
+
1538
+ # make it explicit, finetune by including extra input channels
1539
+ if exists(self.finetune_keys) and k in self.finetune_keys:
1540
+ new_entry = None
1541
+ for name, param in self.named_parameters():
1542
+ if name in self.finetune_keys:
1543
+ print(
1544
+ f"modifying key '{name}' and keeping its original {self.keep_dims} (channels) dimensions only")
1545
+ new_entry = torch.zeros_like(param) # zero init
1546
+ assert exists(new_entry), 'did not find matching parameter to modify'
1547
+ new_entry[:, :self.keep_dims, ...] = sd[k]
1548
+ sd[k] = new_entry
1549
+
1550
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
1551
+ sd, strict=False)
1552
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
1553
+ if len(missing) > 0:
1554
+ print(f"Missing Keys: {missing}")
1555
+ if len(unexpected) > 0:
1556
+ print(f"Unexpected Keys: {unexpected}")
1557
+
1558
+ @torch.no_grad()
1559
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
1560
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
1561
+ plot_diffusion_rows=True, unconditional_guidance_scale=1., unconditional_guidance_label=None,
1562
+ use_ema_scope=True,
1563
+ **kwargs):
1564
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
1565
+ use_ddim = ddim_steps is not None
1566
+
1567
+ log = dict()
1568
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, bs=N, return_first_stage_outputs=True)
1569
+ c_cat, c = c["c_concat"][0], c["c_crossattn"][0]
1570
+ N = min(x.shape[0], N)
1571
+ n_row = min(x.shape[0], n_row)
1572
+ log["inputs"] = x
1573
+ log["reconstruction"] = xrec
1574
+ if self.model.conditioning_key is not None:
1575
+ if hasattr(self.cond_stage_model, "decode"):
1576
+ xc = self.cond_stage_model.decode(c)
1577
+ log["conditioning"] = xc
1578
+ elif self.cond_stage_key in ["caption", "txt"]:
1579
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1580
+ log["conditioning"] = xc
1581
+ elif self.cond_stage_key in ['class_label', 'cls']:
1582
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"], size=x.shape[2] // 25)
1583
+ log['conditioning'] = xc
1584
+ elif isimage(xc):
1585
+ log["conditioning"] = xc
1586
+ if ismap(xc):
1587
+ log["original_conditioning"] = self.to_rgb(xc)
1588
+
1589
+ if not (self.c_concat_log_start is None and self.c_concat_log_end is None):
1590
+ log["c_concat_decoded"] = self.decode_first_stage(c_cat[:, self.c_concat_log_start:self.c_concat_log_end])
1591
+
1592
+ if plot_diffusion_rows:
1593
+ # get diffusion row
1594
+ diffusion_row = list()
1595
+ z_start = z[:n_row]
1596
+ for t in range(self.num_timesteps):
1597
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
1598
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
1599
+ t = t.to(self.device).long()
1600
+ noise = torch.randn_like(z_start)
1601
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
1602
+ diffusion_row.append(self.decode_first_stage(z_noisy))
1603
+
1604
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
1605
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
1606
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
1607
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
1608
+ log["diffusion_row"] = diffusion_grid
1609
+
1610
+ if sample:
1611
+ # get denoise row
1612
+ with ema_scope("Sampling"):
1613
+ samples, z_denoise_row = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
1614
+ batch_size=N, ddim=use_ddim,
1615
+ ddim_steps=ddim_steps, eta=ddim_eta)
1616
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
1617
+ x_samples = self.decode_first_stage(samples)
1618
+ log["samples"] = x_samples
1619
+ if plot_denoise_rows:
1620
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
1621
+ log["denoise_row"] = denoise_grid
1622
+
1623
+ if unconditional_guidance_scale > 1.0:
1624
+ uc_cross = self.get_unconditional_conditioning(N, unconditional_guidance_label)
1625
+ uc_cat = c_cat
1626
+ uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross]}
1627
+ with ema_scope("Sampling with classifier-free guidance"):
1628
+ samples_cfg, _ = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
1629
+ batch_size=N, ddim=use_ddim,
1630
+ ddim_steps=ddim_steps, eta=ddim_eta,
1631
+ unconditional_guidance_scale=unconditional_guidance_scale,
1632
+ unconditional_conditioning=uc_full,
1633
+ )
1634
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1635
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1636
+
1637
+ return log
1638
+
1639
+
1640
+ class LatentInpaintDiffusion(LatentFinetuneDiffusion):
1641
+ """
1642
+ can either run as pure inpainting model (only concat mode) or with mixed conditionings,
1643
+ e.g. mask as concat and text via cross-attn.
1644
+ To disable finetuning mode, set finetune_keys to None
1645
+ """
1646
+
1647
+ def __init__(self,
1648
+ concat_keys=("mask", "masked_image"),
1649
+ masked_image_key="masked_image",
1650
+ *args, **kwargs
1651
+ ):
1652
+ super().__init__(concat_keys, *args, **kwargs)
1653
+ self.masked_image_key = masked_image_key
1654
+ assert self.masked_image_key in concat_keys
1655
+
1656
+ @torch.no_grad()
1657
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1658
+ # note: restricted to non-trainable encoders currently
1659
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for inpainting'
1660
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1661
+ force_c_encode=True, return_original_cond=True, bs=bs)
1662
+
1663
+ assert exists(self.concat_keys)
1664
+ c_cat = list()
1665
+ for ck in self.concat_keys:
1666
+ cc = rearrange(batch[ck], 'b h w c -> b c h w').to(memory_format=torch.contiguous_format).float()
1667
+ if bs is not None:
1668
+ cc = cc[:bs]
1669
+ cc = cc.to(self.device)
1670
+ bchw = z.shape
1671
+ if ck != self.masked_image_key:
1672
+ cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
1673
+ else:
1674
+ cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
1675
+ c_cat.append(cc)
1676
+ c_cat = torch.cat(c_cat, dim=1)
1677
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1678
+ if return_first_stage_outputs:
1679
+ return z, all_conds, x, xrec, xc
1680
+ return z, all_conds
1681
+
1682
+ @torch.no_grad()
1683
+ def log_images(self, *args, **kwargs):
1684
+ log = super(LatentInpaintDiffusion, self).log_images(*args, **kwargs)
1685
+ log["masked_image"] = rearrange(args[0]["masked_image"],
1686
+ 'b h w c -> b c h w').to(memory_format=torch.contiguous_format).float()
1687
+ return log
1688
+
1689
+
1690
+ class LatentDepth2ImageDiffusion(LatentFinetuneDiffusion):
1691
+ """
1692
+ condition on monocular depth estimation
1693
+ """
1694
+
1695
+ def __init__(self, depth_stage_config, concat_keys=("midas_in",), *args, **kwargs):
1696
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
1697
+ self.depth_model = instantiate_from_config(depth_stage_config)
1698
+ self.depth_stage_key = concat_keys[0]
1699
+
1700
+ @torch.no_grad()
1701
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1702
+ # note: restricted to non-trainable encoders currently
1703
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for depth2img'
1704
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1705
+ force_c_encode=True, return_original_cond=True, bs=bs)
1706
+
1707
+ assert exists(self.concat_keys)
1708
+ assert len(self.concat_keys) == 1
1709
+ c_cat = list()
1710
+ for ck in self.concat_keys:
1711
+ cc = batch[ck]
1712
+ if bs is not None:
1713
+ cc = cc[:bs]
1714
+ cc = cc.to(self.device)
1715
+ cc = self.depth_model(cc)
1716
+ cc = torch.nn.functional.interpolate(
1717
+ cc,
1718
+ size=z.shape[2:],
1719
+ mode="bicubic",
1720
+ align_corners=False,
1721
+ )
1722
+
1723
+ depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
1724
+ keepdim=True)
1725
+ cc = 2. * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.
1726
+ c_cat.append(cc)
1727
+ c_cat = torch.cat(c_cat, dim=1)
1728
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1729
+ if return_first_stage_outputs:
1730
+ return z, all_conds, x, xrec, xc
1731
+ return z, all_conds
1732
+
1733
+ @torch.no_grad()
1734
+ def log_images(self, *args, **kwargs):
1735
+ log = super().log_images(*args, **kwargs)
1736
+ depth = self.depth_model(args[0][self.depth_stage_key])
1737
+ depth_min, depth_max = torch.amin(depth, dim=[1, 2, 3], keepdim=True), \
1738
+ torch.amax(depth, dim=[1, 2, 3], keepdim=True)
1739
+ log["depth"] = 2. * (depth - depth_min) / (depth_max - depth_min) - 1.
1740
+ return log
1741
+
1742
+
1743
+ class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion):
1744
+ """
1745
+ condition on low-res image (and optionally on some spatial noise augmentation)
1746
+ """
1747
+ def __init__(self, concat_keys=("lr",), reshuffle_patch_size=None,
1748
+ low_scale_config=None, low_scale_key=None, *args, **kwargs):
1749
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
1750
+ self.reshuffle_patch_size = reshuffle_patch_size
1751
+ self.low_scale_model = None
1752
+ if low_scale_config is not None:
1753
+ print("Initializing a low-scale model")
1754
+ assert exists(low_scale_key)
1755
+ self.instantiate_low_stage(low_scale_config)
1756
+ self.low_scale_key = low_scale_key
1757
+
1758
+ def instantiate_low_stage(self, config):
1759
+ model = instantiate_from_config(config)
1760
+ self.low_scale_model = model.eval()
1761
+ self.low_scale_model.train = disabled_train
1762
+ for param in self.low_scale_model.parameters():
1763
+ param.requires_grad = False
1764
+
1765
+ @torch.no_grad()
1766
+ def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
1767
+ # note: restricted to non-trainable encoders currently
1768
+ assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for upscaling-ft'
1769
+ z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
1770
+ force_c_encode=True, return_original_cond=True, bs=bs)
1771
+
1772
+ assert exists(self.concat_keys)
1773
+ assert len(self.concat_keys) == 1
1774
+ # optionally make spatial noise_level here
1775
+ c_cat = list()
1776
+ noise_level = None
1777
+ for ck in self.concat_keys:
1778
+ cc = batch[ck]
1779
+ cc = rearrange(cc, 'b h w c -> b c h w')
1780
+ if exists(self.reshuffle_patch_size):
1781
+ assert isinstance(self.reshuffle_patch_size, int)
1782
+ cc = rearrange(cc, 'b c (p1 h) (p2 w) -> b (p1 p2 c) h w',
1783
+ p1=self.reshuffle_patch_size, p2=self.reshuffle_patch_size)
1784
+ if bs is not None:
1785
+ cc = cc[:bs]
1786
+ cc = cc.to(self.device)
1787
+ if exists(self.low_scale_model) and ck == self.low_scale_key:
1788
+ cc, noise_level = self.low_scale_model(cc)
1789
+ c_cat.append(cc)
1790
+ c_cat = torch.cat(c_cat, dim=1)
1791
+ if exists(noise_level):
1792
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c], "c_adm": noise_level}
1793
+ else:
1794
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
1795
+ if return_first_stage_outputs:
1796
+ return z, all_conds, x, xrec, xc
1797
+ return z, all_conds
1798
+
1799
+ @torch.no_grad()
1800
+ def log_images(self, *args, **kwargs):
1801
+ log = super().log_images(*args, **kwargs)
1802
+ log["lr"] = rearrange(args[0]["lr"], 'b h w c -> b c h w')
1803
+ return log
1804
+
1805
+
1806
+ class ImageEmbeddingConditionedLatentDiffusion(LatentDiffusion):
1807
+ def __init__(self, embedder_config=None, embedding_key="jpg", embedding_dropout=0.5,
1808
+ freeze_embedder=True, noise_aug_config=None, *args, **kwargs):
1809
+ super().__init__(*args, **kwargs)
1810
+ self.embed_key = embedding_key
1811
+ self.embedding_dropout = embedding_dropout
1812
+ # self._init_embedder(embedder_config, freeze_embedder)
1813
+ self._init_noise_aug(noise_aug_config)
1814
+
1815
+ def _init_embedder(self, config, freeze=True):
1816
+ embedder = instantiate_from_config(config)
1817
+ if freeze:
1818
+ self.embedder = embedder.eval()
1819
+ self.embedder.train = disabled_train
1820
+ for param in self.embedder.parameters():
1821
+ param.requires_grad = False
1822
+
1823
+ def _init_noise_aug(self, config):
1824
+ if config is not None:
1825
+ # use the KARLO schedule for noise augmentation on CLIP image embeddings
1826
+ noise_augmentor = instantiate_from_config(config)
1827
+ assert isinstance(noise_augmentor, nn.Module)
1828
+ noise_augmentor = noise_augmentor.eval()
1829
+ noise_augmentor.train = disabled_train
1830
+ self.noise_augmentor = noise_augmentor
1831
+ else:
1832
+ self.noise_augmentor = None
1833
+
1834
+ def get_input(self, batch, k, cond_key=None, bs=None, **kwargs):
1835
+ outputs = LatentDiffusion.get_input(self, batch, k, bs=bs, **kwargs)
1836
+ z, c = outputs[0], outputs[1]
1837
+ img = batch[self.embed_key][:bs]
1838
+ img = rearrange(img, 'b h w c -> b c h w')
1839
+ c_adm = self.embedder(img)
1840
+ if self.noise_augmentor is not None:
1841
+ c_adm, noise_level_emb = self.noise_augmentor(c_adm)
1842
+ # assume this gives embeddings of noise levels
1843
+ c_adm = torch.cat((c_adm, noise_level_emb), 1)
1844
+ if self.training:
1845
+ c_adm = torch.bernoulli((1. - self.embedding_dropout) * torch.ones(c_adm.shape[0],
1846
+ device=c_adm.device)[:, None]) * c_adm
1847
+ all_conds = {"c_crossattn": [c], "c_adm": c_adm}
1848
+ noutputs = [z, all_conds]
1849
+ noutputs.extend(outputs[2:])
1850
+ return noutputs
1851
+
1852
+ @torch.no_grad()
1853
+ def log_images(self, batch, N=8, n_row=4, **kwargs):
1854
+ log = dict()
1855
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, bs=N, return_first_stage_outputs=True,
1856
+ return_original_cond=True)
1857
+ log["inputs"] = x
1858
+ log["reconstruction"] = xrec
1859
+ assert self.model.conditioning_key is not None
1860
+ assert self.cond_stage_key in ["caption", "txt"]
1861
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch[self.cond_stage_key], size=x.shape[2] // 25)
1862
+ log["conditioning"] = xc
1863
+ uc = self.get_unconditional_conditioning(N, kwargs.get('unconditional_guidance_label', ''))
1864
+ unconditional_guidance_scale = kwargs.get('unconditional_guidance_scale', 5.)
1865
+
1866
+ uc_ = {"c_crossattn": [uc], "c_adm": c["c_adm"]}
1867
+ ema_scope = self.ema_scope if kwargs.get('use_ema_scope', True) else nullcontext
1868
+ with ema_scope(f"Sampling"):
1869
+ samples_cfg, _ = self.sample_log(cond=c, batch_size=N, ddim=True,
1870
+ ddim_steps=kwargs.get('ddim_steps', 50), eta=kwargs.get('ddim_eta', 0.),
1871
+ unconditional_guidance_scale=unconditional_guidance_scale,
1872
+ unconditional_conditioning=uc_, )
1873
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
1874
+ log[f"samplescfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
1875
+ return log
comfy/ldm/models/diffusion/dpm_solver/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .sampler import DPMSolverSampler
comfy/ldm/models/diffusion/dpm_solver/dpm_solver.py ADDED
@@ -0,0 +1,1163 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
19
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
20
+ ***
21
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
22
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
23
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
24
+ log_alpha_t = self.marginal_log_mean_coeff(t)
25
+ sigma_t = self.marginal_std(t)
26
+ lambda_t = self.marginal_lambda(t)
27
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
28
+ t = self.inverse_lambda(lambda_t)
29
+ ===============================================================
30
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
31
+ 1. For discrete-time DPMs:
32
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
33
+ t_i = (i + 1) / N
34
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
35
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
36
+ Args:
37
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
38
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
39
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
40
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
41
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
42
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
43
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
44
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
45
+ and
46
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
47
+ 2. For continuous-time DPMs:
48
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
49
+ schedule are the default settings in DDPM and improved-DDPM:
50
+ Args:
51
+ beta_min: A `float` number. The smallest beta for the linear schedule.
52
+ beta_max: A `float` number. The largest beta for the linear schedule.
53
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
54
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
55
+ T: A `float` number. The ending time of the forward process.
56
+ ===============================================================
57
+ Args:
58
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
59
+ 'linear' or 'cosine' for continuous-time DPMs.
60
+ Returns:
61
+ A wrapper object of the forward SDE (VP type).
62
+
63
+ ===============================================================
64
+ Example:
65
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
66
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
67
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
68
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
69
+ # For continuous-time DPMs (VPSDE), linear schedule:
70
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
71
+ """
72
+
73
+ if schedule not in ['discrete', 'linear', 'cosine']:
74
+ raise ValueError(
75
+ "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
76
+ schedule))
77
+
78
+ self.schedule = schedule
79
+ if schedule == 'discrete':
80
+ if betas is not None:
81
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
82
+ else:
83
+ assert alphas_cumprod is not None
84
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
85
+ self.total_N = len(log_alphas)
86
+ self.T = 1.
87
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
88
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
89
+ else:
90
+ self.total_N = 1000
91
+ self.beta_0 = continuous_beta_0
92
+ self.beta_1 = continuous_beta_1
93
+ self.cosine_s = 0.008
94
+ self.cosine_beta_max = 999.
95
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
96
+ 1. + self.cosine_s) / math.pi - self.cosine_s
97
+ self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
98
+ self.schedule = schedule
99
+ if schedule == 'cosine':
100
+ # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
101
+ # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
102
+ self.T = 0.9946
103
+ else:
104
+ self.T = 1.
105
+
106
+ def marginal_log_mean_coeff(self, t):
107
+ """
108
+ Compute log(alpha_t) of a given continuous-time label t in [0, T].
109
+ """
110
+ if self.schedule == 'discrete':
111
+ return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
112
+ self.log_alpha_array.to(t.device)).reshape((-1))
113
+ elif self.schedule == 'linear':
114
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
115
+ elif self.schedule == 'cosine':
116
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
117
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
118
+ return log_alpha_t
119
+
120
+ def marginal_alpha(self, t):
121
+ """
122
+ Compute alpha_t of a given continuous-time label t in [0, T].
123
+ """
124
+ return torch.exp(self.marginal_log_mean_coeff(t))
125
+
126
+ def marginal_std(self, t):
127
+ """
128
+ Compute sigma_t of a given continuous-time label t in [0, T].
129
+ """
130
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
131
+
132
+ def marginal_lambda(self, t):
133
+ """
134
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
135
+ """
136
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
137
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
138
+ return log_mean_coeff - log_std
139
+
140
+ def inverse_lambda(self, lamb):
141
+ """
142
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
143
+ """
144
+ if self.schedule == 'linear':
145
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
146
+ Delta = self.beta_0 ** 2 + tmp
147
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
148
+ elif self.schedule == 'discrete':
149
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
150
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
151
+ torch.flip(self.t_array.to(lamb.device), [1]))
152
+ return t.reshape((-1,))
153
+ else:
154
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
155
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
156
+ 1. + self.cosine_s) / math.pi - self.cosine_s
157
+ t = t_fn(log_alpha)
158
+ return t
159
+
160
+
161
+ def model_wrapper(
162
+ model,
163
+ noise_schedule,
164
+ model_type="noise",
165
+ model_kwargs={},
166
+ guidance_type="uncond",
167
+ condition=None,
168
+ unconditional_condition=None,
169
+ guidance_scale=1.,
170
+ classifier_fn=None,
171
+ classifier_kwargs={},
172
+ ):
173
+ """Create a wrapper function for the noise prediction model.
174
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
175
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
176
+ We support four types of the diffusion model by setting `model_type`:
177
+ 1. "noise": noise prediction model. (Trained by predicting noise).
178
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
179
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
180
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
181
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
182
+ arXiv preprint arXiv:2202.00512 (2022).
183
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
184
+ arXiv preprint arXiv:2210.02303 (2022).
185
+
186
+ 4. "score": marginal score function. (Trained by denoising score matching).
187
+ Note that the score function and the noise prediction model follows a simple relationship:
188
+ ```
189
+ noise(x_t, t) = -sigma_t * score(x_t, t)
190
+ ```
191
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
192
+ 1. "uncond": unconditional sampling by DPMs.
193
+ The input `model` has the following format:
194
+ ``
195
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
196
+ ``
197
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
198
+ The input `model` has the following format:
199
+ ``
200
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
201
+ ``
202
+ The input `classifier_fn` has the following format:
203
+ ``
204
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
205
+ ``
206
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
207
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
208
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
209
+ The input `model` has the following format:
210
+ ``
211
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
212
+ ``
213
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
214
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
215
+ arXiv preprint arXiv:2207.12598 (2022).
216
+
217
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
218
+ or continuous-time labels (i.e. epsilon to T).
219
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
220
+ ``
221
+ def model_fn(x, t_continuous) -> noise:
222
+ t_input = get_model_input_time(t_continuous)
223
+ return noise_pred(model, x, t_input, **model_kwargs)
224
+ ``
225
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
226
+ ===============================================================
227
+ Args:
228
+ model: A diffusion model with the corresponding format described above.
229
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
230
+ model_type: A `str`. The parameterization type of the diffusion model.
231
+ "noise" or "x_start" or "v" or "score".
232
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
233
+ guidance_type: A `str`. The type of the guidance for sampling.
234
+ "uncond" or "classifier" or "classifier-free".
235
+ condition: A pytorch tensor. The condition for the guided sampling.
236
+ Only used for "classifier" or "classifier-free" guidance type.
237
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
238
+ Only used for "classifier-free" guidance type.
239
+ guidance_scale: A `float`. The scale for the guided sampling.
240
+ classifier_fn: A classifier function. Only used for the classifier guidance.
241
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
242
+ Returns:
243
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
244
+ """
245
+
246
+ def get_model_input_time(t_continuous):
247
+ """
248
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
249
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
250
+ For continuous-time DPMs, we just use `t_continuous`.
251
+ """
252
+ if noise_schedule.schedule == 'discrete':
253
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
254
+ else:
255
+ return t_continuous
256
+
257
+ def noise_pred_fn(x, t_continuous, cond=None):
258
+ if t_continuous.reshape((-1,)).shape[0] == 1:
259
+ t_continuous = t_continuous.expand((x.shape[0]))
260
+ t_input = get_model_input_time(t_continuous)
261
+ if cond is None:
262
+ output = model(x, t_input, **model_kwargs)
263
+ else:
264
+ output = model(x, t_input, cond, **model_kwargs)
265
+ if model_type == "noise":
266
+ return output
267
+ elif model_type == "x_start":
268
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
269
+ dims = x.dim()
270
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
271
+ elif model_type == "v":
272
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
273
+ dims = x.dim()
274
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
275
+ elif model_type == "score":
276
+ sigma_t = noise_schedule.marginal_std(t_continuous)
277
+ dims = x.dim()
278
+ return -expand_dims(sigma_t, dims) * output
279
+
280
+ def cond_grad_fn(x, t_input):
281
+ """
282
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
283
+ """
284
+ with torch.enable_grad():
285
+ x_in = x.detach().requires_grad_(True)
286
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
287
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
288
+
289
+ def model_fn(x, t_continuous):
290
+ """
291
+ The noise predicition model function that is used for DPM-Solver.
292
+ """
293
+ if t_continuous.reshape((-1,)).shape[0] == 1:
294
+ t_continuous = t_continuous.expand((x.shape[0]))
295
+ if guidance_type == "uncond":
296
+ return noise_pred_fn(x, t_continuous)
297
+ elif guidance_type == "classifier":
298
+ assert classifier_fn is not None
299
+ t_input = get_model_input_time(t_continuous)
300
+ cond_grad = cond_grad_fn(x, t_input)
301
+ sigma_t = noise_schedule.marginal_std(t_continuous)
302
+ noise = noise_pred_fn(x, t_continuous)
303
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
304
+ elif guidance_type == "classifier-free":
305
+ if guidance_scale == 1. or unconditional_condition is None:
306
+ return noise_pred_fn(x, t_continuous, cond=condition)
307
+ else:
308
+ x_in = torch.cat([x] * 2)
309
+ t_in = torch.cat([t_continuous] * 2)
310
+ if isinstance(condition, dict):
311
+ assert isinstance(unconditional_condition, dict)
312
+ c_in = dict()
313
+ for k in condition:
314
+ if isinstance(condition[k], list):
315
+ c_in[k] = [torch.cat([unconditional_condition[k][i], condition[k][i]]) for i in range(len(condition[k]))]
316
+ else:
317
+ c_in[k] = torch.cat([unconditional_condition[k], condition[k]])
318
+ else:
319
+ c_in = torch.cat([unconditional_condition, condition])
320
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
321
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
322
+
323
+ assert model_type in ["noise", "x_start", "v"]
324
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
325
+ return model_fn
326
+
327
+
328
+ class DPM_Solver:
329
+ def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
330
+ """Construct a DPM-Solver.
331
+ We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
332
+ If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
333
+ If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
334
+ In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
335
+ The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
336
+ Args:
337
+ model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
338
+ ``
339
+ def model_fn(x, t_continuous):
340
+ return noise
341
+ ``
342
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
343
+ predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
344
+ thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
345
+ max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
346
+
347
+ [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.
348
+ """
349
+ self.model = model_fn
350
+ self.noise_schedule = noise_schedule
351
+ self.predict_x0 = predict_x0
352
+ self.thresholding = thresholding
353
+ self.max_val = max_val
354
+
355
+ def noise_prediction_fn(self, x, t):
356
+ """
357
+ Return the noise prediction model.
358
+ """
359
+ return self.model(x, t)
360
+
361
+ def data_prediction_fn(self, x, t):
362
+ """
363
+ Return the data prediction model (with thresholding).
364
+ """
365
+ noise = self.noise_prediction_fn(x, t)
366
+ dims = x.dim()
367
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
368
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
369
+ if self.thresholding:
370
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
371
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
372
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
373
+ x0 = torch.clamp(x0, -s, s) / s
374
+ return x0
375
+
376
+ def model_fn(self, x, t):
377
+ """
378
+ Convert the model to the noise prediction model or the data prediction model.
379
+ """
380
+ if self.predict_x0:
381
+ return self.data_prediction_fn(x, t)
382
+ else:
383
+ return self.noise_prediction_fn(x, t)
384
+
385
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
386
+ """Compute the intermediate time steps for sampling.
387
+ Args:
388
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
389
+ - 'logSNR': uniform logSNR for the time steps.
390
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
391
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
392
+ t_T: A `float`. The starting time of the sampling (default is T).
393
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
394
+ N: A `int`. The total number of the spacing of the time steps.
395
+ device: A torch device.
396
+ Returns:
397
+ A pytorch tensor of the time steps, with the shape (N + 1,).
398
+ """
399
+ if skip_type == 'logSNR':
400
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
401
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
402
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
403
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
404
+ elif skip_type == 'time_uniform':
405
+ return torch.linspace(t_T, t_0, N + 1).to(device)
406
+ elif skip_type == 'time_quadratic':
407
+ t_order = 2
408
+ t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
409
+ return t
410
+ else:
411
+ raise ValueError(
412
+ "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
413
+
414
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
415
+ """
416
+ Get the order of each step for sampling by the singlestep DPM-Solver.
417
+ We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
418
+ Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
419
+ - If order == 1:
420
+ We take `steps` of DPM-Solver-1 (i.e. DDIM).
421
+ - If order == 2:
422
+ - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
423
+ - If steps % 2 == 0, we use K steps of DPM-Solver-2.
424
+ - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
425
+ - If order == 3:
426
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
427
+ - 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.
428
+ - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
429
+ - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
430
+ ============================================
431
+ Args:
432
+ order: A `int`. The max order for the solver (2 or 3).
433
+ steps: A `int`. The total number of function evaluations (NFE).
434
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
435
+ - 'logSNR': uniform logSNR for the time steps.
436
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
437
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
438
+ t_T: A `float`. The starting time of the sampling (default is T).
439
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
440
+ device: A torch device.
441
+ Returns:
442
+ orders: A list of the solver order of each step.
443
+ """
444
+ if order == 3:
445
+ K = steps // 3 + 1
446
+ if steps % 3 == 0:
447
+ orders = [3, ] * (K - 2) + [2, 1]
448
+ elif steps % 3 == 1:
449
+ orders = [3, ] * (K - 1) + [1]
450
+ else:
451
+ orders = [3, ] * (K - 1) + [2]
452
+ elif order == 2:
453
+ if steps % 2 == 0:
454
+ K = steps // 2
455
+ orders = [2, ] * K
456
+ else:
457
+ K = steps // 2 + 1
458
+ orders = [2, ] * (K - 1) + [1]
459
+ elif order == 1:
460
+ K = 1
461
+ orders = [1, ] * steps
462
+ else:
463
+ raise ValueError("'order' must be '1' or '2' or '3'.")
464
+ if skip_type == 'logSNR':
465
+ # To reproduce the results in DPM-Solver paper
466
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
467
+ else:
468
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
469
+ torch.cumsum(torch.tensor([0, ] + orders)).to(device)]
470
+ return timesteps_outer, orders
471
+
472
+ def denoise_to_zero_fn(self, x, s):
473
+ """
474
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
475
+ """
476
+ return self.data_prediction_fn(x, s)
477
+
478
+ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
479
+ """
480
+ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
481
+ Args:
482
+ x: A pytorch tensor. The initial value at time `s`.
483
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
484
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
485
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
486
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
487
+ return_intermediate: A `bool`. If true, also return the model value at time `s`.
488
+ Returns:
489
+ x_t: A pytorch tensor. The approximated solution at time `t`.
490
+ """
491
+ ns = self.noise_schedule
492
+ dims = x.dim()
493
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
494
+ h = lambda_t - lambda_s
495
+ log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
496
+ sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
497
+ alpha_t = torch.exp(log_alpha_t)
498
+
499
+ if self.predict_x0:
500
+ phi_1 = torch.expm1(-h)
501
+ if model_s is None:
502
+ model_s = self.model_fn(x, s)
503
+ x_t = (
504
+ expand_dims(sigma_t / sigma_s, dims) * x
505
+ - expand_dims(alpha_t * phi_1, dims) * model_s
506
+ )
507
+ if return_intermediate:
508
+ return x_t, {'model_s': model_s}
509
+ else:
510
+ return x_t
511
+ else:
512
+ phi_1 = torch.expm1(h)
513
+ if model_s is None:
514
+ model_s = self.model_fn(x, s)
515
+ x_t = (
516
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
517
+ - expand_dims(sigma_t * phi_1, dims) * model_s
518
+ )
519
+ if return_intermediate:
520
+ return x_t, {'model_s': model_s}
521
+ else:
522
+ return x_t
523
+
524
+ def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
525
+ solver_type='dpm_solver'):
526
+ """
527
+ Singlestep solver DPM-Solver-2 from time `s` to time `t`.
528
+ Args:
529
+ x: A pytorch tensor. The initial value at time `s`.
530
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
531
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
532
+ r1: A `float`. The hyperparameter of the second-order solver.
533
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
534
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
535
+ return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
536
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
537
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
538
+ Returns:
539
+ x_t: A pytorch tensor. The approximated solution at time `t`.
540
+ """
541
+ if solver_type not in ['dpm_solver', 'taylor']:
542
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
543
+ if r1 is None:
544
+ r1 = 0.5
545
+ ns = self.noise_schedule
546
+ dims = x.dim()
547
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
548
+ h = lambda_t - lambda_s
549
+ lambda_s1 = lambda_s + r1 * h
550
+ s1 = ns.inverse_lambda(lambda_s1)
551
+ log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
552
+ s1), ns.marginal_log_mean_coeff(t)
553
+ sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
554
+ alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
555
+
556
+ if self.predict_x0:
557
+ phi_11 = torch.expm1(-r1 * h)
558
+ phi_1 = torch.expm1(-h)
559
+
560
+ if model_s is None:
561
+ model_s = self.model_fn(x, s)
562
+ x_s1 = (
563
+ expand_dims(sigma_s1 / sigma_s, dims) * x
564
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
565
+ )
566
+ model_s1 = self.model_fn(x_s1, s1)
567
+ if solver_type == 'dpm_solver':
568
+ x_t = (
569
+ expand_dims(sigma_t / sigma_s, dims) * x
570
+ - expand_dims(alpha_t * phi_1, dims) * model_s
571
+ - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
572
+ )
573
+ elif solver_type == 'taylor':
574
+ x_t = (
575
+ expand_dims(sigma_t / sigma_s, dims) * x
576
+ - expand_dims(alpha_t * phi_1, dims) * model_s
577
+ + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
578
+ model_s1 - model_s)
579
+ )
580
+ else:
581
+ phi_11 = torch.expm1(r1 * h)
582
+ phi_1 = torch.expm1(h)
583
+
584
+ if model_s is None:
585
+ model_s = self.model_fn(x, s)
586
+ x_s1 = (
587
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
588
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
589
+ )
590
+ model_s1 = self.model_fn(x_s1, s1)
591
+ if solver_type == 'dpm_solver':
592
+ x_t = (
593
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
594
+ - expand_dims(sigma_t * phi_1, dims) * model_s
595
+ - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
596
+ )
597
+ elif solver_type == 'taylor':
598
+ x_t = (
599
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
600
+ - expand_dims(sigma_t * phi_1, dims) * model_s
601
+ - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
602
+ )
603
+ if return_intermediate:
604
+ return x_t, {'model_s': model_s, 'model_s1': model_s1}
605
+ else:
606
+ return x_t
607
+
608
+ def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
609
+ return_intermediate=False, solver_type='dpm_solver'):
610
+ """
611
+ Singlestep solver DPM-Solver-3 from time `s` to time `t`.
612
+ Args:
613
+ x: A pytorch tensor. The initial value at time `s`.
614
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
615
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
616
+ r1: A `float`. The hyperparameter of the third-order solver.
617
+ r2: A `float`. The hyperparameter of the third-order solver.
618
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
619
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
620
+ model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
621
+ If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
622
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
623
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
624
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
625
+ Returns:
626
+ x_t: A pytorch tensor. The approximated solution at time `t`.
627
+ """
628
+ if solver_type not in ['dpm_solver', 'taylor']:
629
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
630
+ if r1 is None:
631
+ r1 = 1. / 3.
632
+ if r2 is None:
633
+ r2 = 2. / 3.
634
+ ns = self.noise_schedule
635
+ dims = x.dim()
636
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
637
+ h = lambda_t - lambda_s
638
+ lambda_s1 = lambda_s + r1 * h
639
+ lambda_s2 = lambda_s + r2 * h
640
+ s1 = ns.inverse_lambda(lambda_s1)
641
+ s2 = ns.inverse_lambda(lambda_s2)
642
+ log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
643
+ s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
644
+ sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
645
+ s2), ns.marginal_std(t)
646
+ alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
647
+
648
+ if self.predict_x0:
649
+ phi_11 = torch.expm1(-r1 * h)
650
+ phi_12 = torch.expm1(-r2 * h)
651
+ phi_1 = torch.expm1(-h)
652
+ phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
653
+ phi_2 = phi_1 / h + 1.
654
+ phi_3 = phi_2 / h - 0.5
655
+
656
+ if model_s is None:
657
+ model_s = self.model_fn(x, s)
658
+ if model_s1 is None:
659
+ x_s1 = (
660
+ expand_dims(sigma_s1 / sigma_s, dims) * x
661
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
662
+ )
663
+ model_s1 = self.model_fn(x_s1, s1)
664
+ x_s2 = (
665
+ expand_dims(sigma_s2 / sigma_s, dims) * x
666
+ - expand_dims(alpha_s2 * phi_12, dims) * model_s
667
+ + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
668
+ )
669
+ model_s2 = self.model_fn(x_s2, s2)
670
+ if solver_type == 'dpm_solver':
671
+ x_t = (
672
+ expand_dims(sigma_t / sigma_s, dims) * x
673
+ - expand_dims(alpha_t * phi_1, dims) * model_s
674
+ + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
675
+ )
676
+ elif solver_type == 'taylor':
677
+ D1_0 = (1. / r1) * (model_s1 - model_s)
678
+ D1_1 = (1. / r2) * (model_s2 - model_s)
679
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
680
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
681
+ x_t = (
682
+ expand_dims(sigma_t / sigma_s, dims) * x
683
+ - expand_dims(alpha_t * phi_1, dims) * model_s
684
+ + expand_dims(alpha_t * phi_2, dims) * D1
685
+ - expand_dims(alpha_t * phi_3, dims) * D2
686
+ )
687
+ else:
688
+ phi_11 = torch.expm1(r1 * h)
689
+ phi_12 = torch.expm1(r2 * h)
690
+ phi_1 = torch.expm1(h)
691
+ phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
692
+ phi_2 = phi_1 / h - 1.
693
+ phi_3 = phi_2 / h - 0.5
694
+
695
+ if model_s is None:
696
+ model_s = self.model_fn(x, s)
697
+ if model_s1 is None:
698
+ x_s1 = (
699
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
700
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
701
+ )
702
+ model_s1 = self.model_fn(x_s1, s1)
703
+ x_s2 = (
704
+ expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
705
+ - expand_dims(sigma_s2 * phi_12, dims) * model_s
706
+ - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
707
+ )
708
+ model_s2 = self.model_fn(x_s2, s2)
709
+ if solver_type == 'dpm_solver':
710
+ x_t = (
711
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
712
+ - expand_dims(sigma_t * phi_1, dims) * model_s
713
+ - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
714
+ )
715
+ elif solver_type == 'taylor':
716
+ D1_0 = (1. / r1) * (model_s1 - model_s)
717
+ D1_1 = (1. / r2) * (model_s2 - model_s)
718
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
719
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
720
+ x_t = (
721
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
722
+ - expand_dims(sigma_t * phi_1, dims) * model_s
723
+ - expand_dims(sigma_t * phi_2, dims) * D1
724
+ - expand_dims(sigma_t * phi_3, dims) * D2
725
+ )
726
+
727
+ if return_intermediate:
728
+ return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
729
+ else:
730
+ return x_t
731
+
732
+ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
733
+ """
734
+ Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
735
+ Args:
736
+ x: A pytorch tensor. The initial value at time `s`.
737
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
738
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
739
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
740
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
741
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
742
+ Returns:
743
+ x_t: A pytorch tensor. The approximated solution at time `t`.
744
+ """
745
+ if solver_type not in ['dpm_solver', 'taylor']:
746
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
747
+ ns = self.noise_schedule
748
+ dims = x.dim()
749
+ model_prev_1, model_prev_0 = model_prev_list
750
+ t_prev_1, t_prev_0 = t_prev_list
751
+ lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
752
+ t_prev_0), ns.marginal_lambda(t)
753
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
754
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
755
+ alpha_t = torch.exp(log_alpha_t)
756
+
757
+ h_0 = lambda_prev_0 - lambda_prev_1
758
+ h = lambda_t - lambda_prev_0
759
+ r0 = h_0 / h
760
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
761
+ if self.predict_x0:
762
+ if solver_type == 'dpm_solver':
763
+ x_t = (
764
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
765
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
766
+ - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
767
+ )
768
+ elif solver_type == 'taylor':
769
+ x_t = (
770
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
771
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
772
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
773
+ )
774
+ else:
775
+ if solver_type == 'dpm_solver':
776
+ x_t = (
777
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
778
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
779
+ - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
780
+ )
781
+ elif solver_type == 'taylor':
782
+ x_t = (
783
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
784
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
785
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
786
+ )
787
+ return x_t
788
+
789
+ def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
790
+ """
791
+ Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
792
+ Args:
793
+ x: A pytorch tensor. The initial value at time `s`.
794
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
795
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
796
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
797
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
798
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
799
+ Returns:
800
+ x_t: A pytorch tensor. The approximated solution at time `t`.
801
+ """
802
+ ns = self.noise_schedule
803
+ dims = x.dim()
804
+ model_prev_2, model_prev_1, model_prev_0 = model_prev_list
805
+ t_prev_2, t_prev_1, t_prev_0 = t_prev_list
806
+ lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
807
+ t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
808
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
809
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
810
+ alpha_t = torch.exp(log_alpha_t)
811
+
812
+ h_1 = lambda_prev_1 - lambda_prev_2
813
+ h_0 = lambda_prev_0 - lambda_prev_1
814
+ h = lambda_t - lambda_prev_0
815
+ r0, r1 = h_0 / h, h_1 / h
816
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
817
+ D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
818
+ D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
819
+ D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
820
+ if self.predict_x0:
821
+ x_t = (
822
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
823
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
824
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
825
+ - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
826
+ )
827
+ else:
828
+ x_t = (
829
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
830
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
831
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
832
+ - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
833
+ )
834
+ return x_t
835
+
836
+ def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
837
+ r2=None):
838
+ """
839
+ Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
840
+ Args:
841
+ x: A pytorch tensor. The initial value at time `s`.
842
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
843
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
844
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
845
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
846
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
847
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
848
+ r1: A `float`. The hyperparameter of the second-order or third-order solver.
849
+ r2: A `float`. The hyperparameter of the third-order solver.
850
+ Returns:
851
+ x_t: A pytorch tensor. The approximated solution at time `t`.
852
+ """
853
+ if order == 1:
854
+ return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
855
+ elif order == 2:
856
+ return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
857
+ solver_type=solver_type, r1=r1)
858
+ elif order == 3:
859
+ return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
860
+ solver_type=solver_type, r1=r1, r2=r2)
861
+ else:
862
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
863
+
864
+ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
865
+ """
866
+ Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
867
+ Args:
868
+ x: A pytorch tensor. The initial value at time `s`.
869
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
870
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
871
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
872
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
873
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
874
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
875
+ Returns:
876
+ x_t: A pytorch tensor. The approximated solution at time `t`.
877
+ """
878
+ if order == 1:
879
+ return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
880
+ elif order == 2:
881
+ return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
882
+ elif order == 3:
883
+ return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
884
+ else:
885
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
886
+
887
+ 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,
888
+ solver_type='dpm_solver'):
889
+ """
890
+ The adaptive step size solver based on singlestep DPM-Solver.
891
+ Args:
892
+ x: A pytorch tensor. The initial value at time `t_T`.
893
+ order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
894
+ t_T: A `float`. The starting time of the sampling (default is T).
895
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
896
+ h_init: A `float`. The initial step size (for logSNR).
897
+ atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
898
+ rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
899
+ theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
900
+ t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
901
+ current time and `t_0` is less than `t_err`. The default setting is 1e-5.
902
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
903
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
904
+ Returns:
905
+ x_0: A pytorch tensor. The approximated solution at time `t_0`.
906
+ [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.
907
+ """
908
+ ns = self.noise_schedule
909
+ s = t_T * torch.ones((x.shape[0],)).to(x)
910
+ lambda_s = ns.marginal_lambda(s)
911
+ lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
912
+ h = h_init * torch.ones_like(s).to(x)
913
+ x_prev = x
914
+ nfe = 0
915
+ if order == 2:
916
+ r1 = 0.5
917
+ lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
918
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
919
+ solver_type=solver_type,
920
+ **kwargs)
921
+ elif order == 3:
922
+ r1, r2 = 1. / 3., 2. / 3.
923
+ lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
924
+ return_intermediate=True,
925
+ solver_type=solver_type)
926
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
927
+ solver_type=solver_type,
928
+ **kwargs)
929
+ else:
930
+ raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
931
+ while torch.abs((s - t_0)).mean() > t_err:
932
+ t = ns.inverse_lambda(lambda_s + h)
933
+ x_lower, lower_noise_kwargs = lower_update(x, s, t)
934
+ x_higher = higher_update(x, s, t, **lower_noise_kwargs)
935
+ delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
936
+ norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
937
+ E = norm_fn((x_higher - x_lower) / delta).max()
938
+ if torch.all(E <= 1.):
939
+ x = x_higher
940
+ s = t
941
+ x_prev = x_lower
942
+ lambda_s = ns.marginal_lambda(s)
943
+ h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
944
+ nfe += order
945
+ print('adaptive solver nfe', nfe)
946
+ return x
947
+
948
+ def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
949
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
950
+ atol=0.0078, rtol=0.05,
951
+ ):
952
+ """
953
+ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
954
+ =====================================================
955
+ We support the following algorithms for both noise prediction model and data prediction model:
956
+ - 'singlestep':
957
+ Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
958
+ We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
959
+ The total number of function evaluations (NFE) == `steps`.
960
+ Given a fixed NFE == `steps`, the sampling procedure is:
961
+ - If `order` == 1:
962
+ - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
963
+ - If `order` == 2:
964
+ - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
965
+ - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
966
+ - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
967
+ - If `order` == 3:
968
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
969
+ - 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.
970
+ - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
971
+ - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
972
+ - 'multistep':
973
+ Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
974
+ We initialize the first `order` values by lower order multistep solvers.
975
+ Given a fixed NFE == `steps`, the sampling procedure is:
976
+ Denote K = steps.
977
+ - If `order` == 1:
978
+ - We use K steps of DPM-Solver-1 (i.e. DDIM).
979
+ - If `order` == 2:
980
+ - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
981
+ - If `order` == 3:
982
+ - 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.
983
+ - 'singlestep_fixed':
984
+ Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
985
+ We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
986
+ - 'adaptive':
987
+ Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
988
+ We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
989
+ You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
990
+ (NFE) and the sample quality.
991
+ - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
992
+ - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
993
+ =====================================================
994
+ Some advices for choosing the algorithm:
995
+ - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
996
+ Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
997
+ e.g.
998
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
999
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
1000
+ skip_type='time_uniform', method='singlestep')
1001
+ - For **guided sampling with large guidance scale** by DPMs:
1002
+ Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
1003
+ e.g.
1004
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
1005
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
1006
+ skip_type='time_uniform', method='multistep')
1007
+ We support three types of `skip_type`:
1008
+ - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
1009
+ - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
1010
+ - 'time_quadratic': quadratic time for the time steps.
1011
+ =====================================================
1012
+ Args:
1013
+ x: A pytorch tensor. The initial value at time `t_start`
1014
+ e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
1015
+ steps: A `int`. The total number of function evaluations (NFE).
1016
+ t_start: A `float`. The starting time of the sampling.
1017
+ If `T` is None, we use self.noise_schedule.T (default is 1.0).
1018
+ t_end: A `float`. The ending time of the sampling.
1019
+ If `t_end` is None, we use 1. / self.noise_schedule.total_N.
1020
+ e.g. if total_N == 1000, we have `t_end` == 1e-3.
1021
+ For discrete-time DPMs:
1022
+ - We recommend `t_end` == 1. / self.noise_schedule.total_N.
1023
+ For continuous-time DPMs:
1024
+ - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
1025
+ order: A `int`. The order of DPM-Solver.
1026
+ skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
1027
+ method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
1028
+ denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
1029
+ Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
1030
+ This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
1031
+ score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
1032
+ for diffusion models sampling by diffusion SDEs for low-resolutional images
1033
+ (such as CIFAR-10). However, we observed that such trick does not matter for
1034
+ high-resolutional images. As it needs an additional NFE, we do not recommend
1035
+ it for high-resolutional images.
1036
+ lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
1037
+ Only valid for `method=multistep` and `steps < 15`. We empirically find that
1038
+ this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
1039
+ (especially for steps <= 10). So we recommend to set it to be `True`.
1040
+ solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
1041
+ atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1042
+ rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
1043
+ Returns:
1044
+ x_end: A pytorch tensor. The approximated solution at time `t_end`.
1045
+ """
1046
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
1047
+ t_T = self.noise_schedule.T if t_start is None else t_start
1048
+ device = x.device
1049
+ if method == 'adaptive':
1050
+ with torch.no_grad():
1051
+ x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
1052
+ solver_type=solver_type)
1053
+ elif method == 'multistep':
1054
+ assert steps >= order
1055
+ timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1056
+ assert timesteps.shape[0] - 1 == steps
1057
+ with torch.no_grad():
1058
+ vec_t = timesteps[0].expand((x.shape[0]))
1059
+ model_prev_list = [self.model_fn(x, vec_t)]
1060
+ t_prev_list = [vec_t]
1061
+ # Init the first `order` values by lower order multistep DPM-Solver.
1062
+ for init_order in tqdm(range(1, order), desc="DPM init order"):
1063
+ vec_t = timesteps[init_order].expand(x.shape[0])
1064
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
1065
+ solver_type=solver_type)
1066
+ model_prev_list.append(self.model_fn(x, vec_t))
1067
+ t_prev_list.append(vec_t)
1068
+ # Compute the remaining values by `order`-th order multistep DPM-Solver.
1069
+ for step in tqdm(range(order, steps + 1), desc="DPM multistep"):
1070
+ vec_t = timesteps[step].expand(x.shape[0])
1071
+ if lower_order_final and steps < 15:
1072
+ step_order = min(order, steps + 1 - step)
1073
+ else:
1074
+ step_order = order
1075
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order,
1076
+ solver_type=solver_type)
1077
+ for i in range(order - 1):
1078
+ t_prev_list[i] = t_prev_list[i + 1]
1079
+ model_prev_list[i] = model_prev_list[i + 1]
1080
+ t_prev_list[-1] = vec_t
1081
+ # We do not need to evaluate the final model value.
1082
+ if step < steps:
1083
+ model_prev_list[-1] = self.model_fn(x, vec_t)
1084
+ elif method in ['singlestep', 'singlestep_fixed']:
1085
+ if method == 'singlestep':
1086
+ timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
1087
+ skip_type=skip_type,
1088
+ t_T=t_T, t_0=t_0,
1089
+ device=device)
1090
+ elif method == 'singlestep_fixed':
1091
+ K = steps // order
1092
+ orders = [order, ] * K
1093
+ timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
1094
+ for i, order in enumerate(orders):
1095
+ t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
1096
+ timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
1097
+ N=order, device=device)
1098
+ lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
1099
+ vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])
1100
+ h = lambda_inner[-1] - lambda_inner[0]
1101
+ r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
1102
+ r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
1103
+ x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
1104
+ if denoise_to_zero:
1105
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
1106
+ return x
1107
+
1108
+
1109
+ #############################################################
1110
+ # other utility functions
1111
+ #############################################################
1112
+
1113
+ def interpolate_fn(x, xp, yp):
1114
+ """
1115
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
1116
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
1117
+ 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.)
1118
+ Args:
1119
+ 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).
1120
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
1121
+ yp: PyTorch tensor with shape [C, K].
1122
+ Returns:
1123
+ The function values f(x), with shape [N, C].
1124
+ """
1125
+ N, K = x.shape[0], xp.shape[1]
1126
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
1127
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
1128
+ x_idx = torch.argmin(x_indices, dim=2)
1129
+ cand_start_idx = x_idx - 1
1130
+ start_idx = torch.where(
1131
+ torch.eq(x_idx, 0),
1132
+ torch.tensor(1, device=x.device),
1133
+ torch.where(
1134
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1135
+ ),
1136
+ )
1137
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
1138
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
1139
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
1140
+ start_idx2 = torch.where(
1141
+ torch.eq(x_idx, 0),
1142
+ torch.tensor(0, device=x.device),
1143
+ torch.where(
1144
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
1145
+ ),
1146
+ )
1147
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
1148
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
1149
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
1150
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
1151
+ return cand
1152
+
1153
+
1154
+ def expand_dims(v, dims):
1155
+ """
1156
+ Expand the tensor `v` to the dim `dims`.
1157
+ Args:
1158
+ `v`: a PyTorch tensor with shape [N].
1159
+ `dim`: a `int`.
1160
+ Returns:
1161
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
1162
+ """
1163
+ return v[(...,) + (None,) * (dims - 1)]
comfy/ldm/models/diffusion/dpm_solver/sampler.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+ import torch
3
+
4
+ from .dpm_solver import NoiseScheduleVP, model_wrapper, DPM_Solver
5
+
6
+ MODEL_TYPES = {
7
+ "eps": "noise",
8
+ "v": "v"
9
+ }
10
+
11
+
12
+ class DPMSolverSampler(object):
13
+ def __init__(self, model, device=torch.device("cuda"), **kwargs):
14
+ super().__init__()
15
+ self.model = model
16
+ self.device = device
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 != self.device:
23
+ attr = attr.to(self.device)
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
+ ctmp = conditioning[list(conditioning.keys())[0]]
54
+ while isinstance(ctmp, list): ctmp = ctmp[0]
55
+ if isinstance(ctmp, torch.Tensor):
56
+ cbs = ctmp.shape[0]
57
+ if cbs != batch_size:
58
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
59
+ elif isinstance(conditioning, list):
60
+ for ctmp in conditioning:
61
+ if ctmp.shape[0] != batch_size:
62
+ print(f"Warning: Got {ctmp.shape[0]} conditionings but batch-size is {batch_size}")
63
+ else:
64
+ if isinstance(conditioning, torch.Tensor):
65
+ if conditioning.shape[0] != batch_size:
66
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
67
+
68
+ # sampling
69
+ C, H, W = shape
70
+ size = (batch_size, C, H, W)
71
+
72
+ print(f'Data shape for DPM-Solver sampling is {size}, sampling steps {S}')
73
+
74
+ device = self.model.betas.device
75
+ if x_T is None:
76
+ img = torch.randn(size, device=device)
77
+ else:
78
+ img = x_T
79
+
80
+ ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
81
+
82
+ model_fn = model_wrapper(
83
+ lambda x, t, c: self.model.apply_model(x, t, c),
84
+ ns,
85
+ model_type=MODEL_TYPES[self.model.parameterization],
86
+ guidance_type="classifier-free",
87
+ condition=conditioning,
88
+ unconditional_condition=unconditional_conditioning,
89
+ guidance_scale=unconditional_guidance_scale,
90
+ )
91
+
92
+ dpm_solver = DPM_Solver(model_fn, ns, predict_x0=True, thresholding=False)
93
+ x = dpm_solver.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=2,
94
+ lower_order_final=True)
95
+
96
+ return x.to(device), None
comfy/ldm/models/diffusion/plms.py ADDED
@@ -0,0 +1,245 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+ from functools import partial
7
+
8
+ from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
9
+ from ldm.models.diffusion.sampling_util import norm_thresholding
10
+
11
+
12
+ class PLMSSampler(object):
13
+ def __init__(self, model, schedule="linear", device=torch.device("cuda"), **kwargs):
14
+ super().__init__()
15
+ self.model = model
16
+ self.ddpm_num_timesteps = model.num_timesteps
17
+ self.schedule = schedule
18
+ self.device = device
19
+
20
+ def register_buffer(self, name, attr):
21
+ if type(attr) == torch.Tensor:
22
+ if attr.device != self.device:
23
+ attr = attr.to(self.device)
24
+ setattr(self, name, attr)
25
+
26
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
27
+ if ddim_eta != 0:
28
+ raise ValueError('ddim_eta must be 0 for PLMS')
29
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
30
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
31
+ alphas_cumprod = self.model.alphas_cumprod
32
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
33
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
34
+
35
+ self.register_buffer('betas', to_torch(self.model.betas))
36
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
37
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
38
+
39
+ # calculations for diffusion q(x_t | x_{t-1}) and others
40
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
41
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
42
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
43
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
44
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
45
+
46
+ # ddim sampling parameters
47
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
48
+ ddim_timesteps=self.ddim_timesteps,
49
+ eta=ddim_eta,verbose=verbose)
50
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
51
+ self.register_buffer('ddim_alphas', ddim_alphas)
52
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
53
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
54
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
55
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
56
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
57
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
58
+
59
+ @torch.no_grad()
60
+ def sample(self,
61
+ S,
62
+ batch_size,
63
+ shape,
64
+ conditioning=None,
65
+ callback=None,
66
+ normals_sequence=None,
67
+ img_callback=None,
68
+ quantize_x0=False,
69
+ eta=0.,
70
+ mask=None,
71
+ x0=None,
72
+ temperature=1.,
73
+ noise_dropout=0.,
74
+ score_corrector=None,
75
+ corrector_kwargs=None,
76
+ verbose=True,
77
+ x_T=None,
78
+ log_every_t=100,
79
+ unconditional_guidance_scale=1.,
80
+ unconditional_conditioning=None,
81
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
82
+ dynamic_threshold=None,
83
+ **kwargs
84
+ ):
85
+ if conditioning is not None:
86
+ if isinstance(conditioning, dict):
87
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
88
+ if cbs != batch_size:
89
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
90
+ else:
91
+ if conditioning.shape[0] != batch_size:
92
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
93
+
94
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
95
+ # sampling
96
+ C, H, W = shape
97
+ size = (batch_size, C, H, W)
98
+ print(f'Data shape for PLMS sampling is {size}')
99
+
100
+ samples, intermediates = self.plms_sampling(conditioning, size,
101
+ callback=callback,
102
+ img_callback=img_callback,
103
+ quantize_denoised=quantize_x0,
104
+ mask=mask, x0=x0,
105
+ ddim_use_original_steps=False,
106
+ noise_dropout=noise_dropout,
107
+ temperature=temperature,
108
+ score_corrector=score_corrector,
109
+ corrector_kwargs=corrector_kwargs,
110
+ x_T=x_T,
111
+ log_every_t=log_every_t,
112
+ unconditional_guidance_scale=unconditional_guidance_scale,
113
+ unconditional_conditioning=unconditional_conditioning,
114
+ dynamic_threshold=dynamic_threshold,
115
+ )
116
+ return samples, intermediates
117
+
118
+ @torch.no_grad()
119
+ def plms_sampling(self, cond, shape,
120
+ x_T=None, ddim_use_original_steps=False,
121
+ callback=None, timesteps=None, quantize_denoised=False,
122
+ mask=None, x0=None, img_callback=None, log_every_t=100,
123
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
124
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
125
+ dynamic_threshold=None):
126
+ device = self.model.betas.device
127
+ b = shape[0]
128
+ if x_T is None:
129
+ img = torch.randn(shape, device=device)
130
+ else:
131
+ img = x_T
132
+
133
+ if timesteps is None:
134
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
135
+ elif timesteps is not None and not ddim_use_original_steps:
136
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
137
+ timesteps = self.ddim_timesteps[:subset_end]
138
+
139
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
140
+ time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
141
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
142
+ print(f"Running PLMS Sampling with {total_steps} timesteps")
143
+
144
+ iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
145
+ old_eps = []
146
+
147
+ for i, step in enumerate(iterator):
148
+ index = total_steps - i - 1
149
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
150
+ ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], 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_plms(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
+ old_eps=old_eps, t_next=ts_next,
164
+ dynamic_threshold=dynamic_threshold)
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
+ dynamic_threshold=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)
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).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
+ if dynamic_threshold is not None:
218
+ pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
219
+ # direction pointing to x_t
220
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
221
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
222
+ if noise_dropout > 0.:
223
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
224
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
225
+ return x_prev, pred_x0
226
+
227
+ e_t = get_model_output(x, t)
228
+ if len(old_eps) == 0:
229
+ # Pseudo Improved Euler (2nd order)
230
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
231
+ e_t_next = get_model_output(x_prev, t_next)
232
+ e_t_prime = (e_t + e_t_next) / 2
233
+ elif len(old_eps) == 1:
234
+ # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
235
+ e_t_prime = (3 * e_t - old_eps[-1]) / 2
236
+ elif len(old_eps) == 2:
237
+ # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
238
+ e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
239
+ elif len(old_eps) >= 3:
240
+ # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
241
+ e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
242
+
243
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
244
+
245
+ return x_prev, pred_x0, e_t
comfy/ldm/models/diffusion/sampling_util.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+
4
+
5
+ def append_dims(x, target_dims):
6
+ """Appends dimensions to the end of a tensor until it has target_dims dimensions.
7
+ From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py"""
8
+ dims_to_append = target_dims - x.ndim
9
+ if dims_to_append < 0:
10
+ raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
11
+ return x[(...,) + (None,) * dims_to_append]
12
+
13
+
14
+ def norm_thresholding(x0, value):
15
+ s = append_dims(x0.pow(2).flatten(1).mean(1).sqrt().clamp(min=value), x0.ndim)
16
+ return x0 * (value / s)
17
+
18
+
19
+ def spatial_norm_thresholding(x0, value):
20
+ # b c h w
21
+ s = x0.pow(2).mean(1, keepdim=True).sqrt().clamp(min=value)
22
+ return x0 * (value / s)
comfy/ldm/modules/attention.py ADDED
@@ -0,0 +1,591 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from .sub_quadratic_attention import efficient_dot_product_attention
11
+
12
+ import model_management
13
+
14
+ from . import tomesd
15
+
16
+ if model_management.xformers_enabled():
17
+ import xformers
18
+ import xformers.ops
19
+
20
+ # CrossAttn precision handling
21
+ import os
22
+ _ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
23
+
24
+ def exists(val):
25
+ return val is not None
26
+
27
+
28
+ def uniq(arr):
29
+ return{el: True for el in arr}.keys()
30
+
31
+
32
+ def default(val, d):
33
+ if exists(val):
34
+ return val
35
+ return d() if isfunction(d) else d
36
+
37
+
38
+ def max_neg_value(t):
39
+ return -torch.finfo(t.dtype).max
40
+
41
+
42
+ def init_(tensor):
43
+ dim = tensor.shape[-1]
44
+ std = 1 / math.sqrt(dim)
45
+ tensor.uniform_(-std, std)
46
+ return tensor
47
+
48
+
49
+ # feedforward
50
+ class GEGLU(nn.Module):
51
+ def __init__(self, dim_in, dim_out):
52
+ super().__init__()
53
+ self.proj = nn.Linear(dim_in, dim_out * 2)
54
+
55
+ def forward(self, x):
56
+ x, gate = self.proj(x).chunk(2, dim=-1)
57
+ return x * F.gelu(gate)
58
+
59
+
60
+ class FeedForward(nn.Module):
61
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
62
+ super().__init__()
63
+ inner_dim = int(dim * mult)
64
+ dim_out = default(dim_out, dim)
65
+ project_in = nn.Sequential(
66
+ nn.Linear(dim, inner_dim),
67
+ nn.GELU()
68
+ ) if not glu else GEGLU(dim, inner_dim)
69
+
70
+ self.net = nn.Sequential(
71
+ project_in,
72
+ nn.Dropout(dropout),
73
+ nn.Linear(inner_dim, dim_out)
74
+ )
75
+
76
+ def forward(self, x):
77
+ return self.net(x)
78
+
79
+
80
+ def zero_module(module):
81
+ """
82
+ Zero out the parameters of a module and return it.
83
+ """
84
+ for p in module.parameters():
85
+ p.detach().zero_()
86
+ return module
87
+
88
+
89
+ def Normalize(in_channels):
90
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
91
+
92
+
93
+ class SpatialSelfAttention(nn.Module):
94
+ def __init__(self, in_channels):
95
+ super().__init__()
96
+ self.in_channels = in_channels
97
+
98
+ self.norm = Normalize(in_channels)
99
+ self.q = torch.nn.Conv2d(in_channels,
100
+ in_channels,
101
+ kernel_size=1,
102
+ stride=1,
103
+ padding=0)
104
+ self.k = torch.nn.Conv2d(in_channels,
105
+ in_channels,
106
+ kernel_size=1,
107
+ stride=1,
108
+ padding=0)
109
+ self.v = torch.nn.Conv2d(in_channels,
110
+ in_channels,
111
+ kernel_size=1,
112
+ stride=1,
113
+ padding=0)
114
+ self.proj_out = torch.nn.Conv2d(in_channels,
115
+ in_channels,
116
+ kernel_size=1,
117
+ stride=1,
118
+ padding=0)
119
+
120
+ def forward(self, x):
121
+ h_ = x
122
+ h_ = self.norm(h_)
123
+ q = self.q(h_)
124
+ k = self.k(h_)
125
+ v = self.v(h_)
126
+
127
+ # compute attention
128
+ b,c,h,w = q.shape
129
+ q = rearrange(q, 'b c h w -> b (h w) c')
130
+ k = rearrange(k, 'b c h w -> b c (h w)')
131
+ w_ = torch.einsum('bij,bjk->bik', q, k)
132
+
133
+ w_ = w_ * (int(c)**(-0.5))
134
+ w_ = torch.nn.functional.softmax(w_, dim=2)
135
+
136
+ # attend to values
137
+ v = rearrange(v, 'b c h w -> b c (h w)')
138
+ w_ = rearrange(w_, 'b i j -> b j i')
139
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
140
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
141
+ h_ = self.proj_out(h_)
142
+
143
+ return x+h_
144
+
145
+
146
+ class CrossAttentionBirchSan(nn.Module):
147
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
148
+ super().__init__()
149
+ inner_dim = dim_head * heads
150
+ context_dim = default(context_dim, query_dim)
151
+
152
+ self.scale = dim_head ** -0.5
153
+ self.heads = heads
154
+
155
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
156
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
157
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
158
+
159
+ self.to_out = nn.Sequential(
160
+ nn.Linear(inner_dim, query_dim),
161
+ nn.Dropout(dropout)
162
+ )
163
+
164
+ def forward(self, x, context=None, mask=None):
165
+ h = self.heads
166
+
167
+ query = self.to_q(x)
168
+ context = default(context, x)
169
+ key = self.to_k(context)
170
+ value = self.to_v(context)
171
+ del context, x
172
+
173
+ query = query.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1)
174
+ key_t = key.transpose(1,2).unflatten(1, (self.heads, -1)).flatten(end_dim=1)
175
+ del key
176
+ value = value.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1)
177
+
178
+ dtype = query.dtype
179
+ upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32
180
+ if upcast_attention:
181
+ bytes_per_token = torch.finfo(torch.float32).bits//8
182
+ else:
183
+ bytes_per_token = torch.finfo(query.dtype).bits//8
184
+ batch_x_heads, q_tokens, _ = query.shape
185
+ _, _, k_tokens = key_t.shape
186
+ qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
187
+
188
+ mem_free_total, mem_free_torch = model_management.get_free_memory(query.device, True)
189
+
190
+ chunk_threshold_bytes = mem_free_torch * 0.5 #Using only this seems to work better on AMD
191
+
192
+ kv_chunk_size_min = None
193
+
194
+ #not sure at all about the math here
195
+ #TODO: tweak this
196
+ if mem_free_total > 8192 * 1024 * 1024 * 1.3:
197
+ query_chunk_size_x = 1024 * 4
198
+ elif mem_free_total > 4096 * 1024 * 1024 * 1.3:
199
+ query_chunk_size_x = 1024 * 2
200
+ else:
201
+ query_chunk_size_x = 1024
202
+ kv_chunk_size_min_x = None
203
+ kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024
204
+ if kv_chunk_size_x < 1024:
205
+ kv_chunk_size_x = None
206
+
207
+ if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes:
208
+ # the big matmul fits into our memory limit; do everything in 1 chunk,
209
+ # i.e. send it down the unchunked fast-path
210
+ query_chunk_size = q_tokens
211
+ kv_chunk_size = k_tokens
212
+ else:
213
+ query_chunk_size = query_chunk_size_x
214
+ kv_chunk_size = kv_chunk_size_x
215
+ kv_chunk_size_min = kv_chunk_size_min_x
216
+
217
+ hidden_states = efficient_dot_product_attention(
218
+ query,
219
+ key_t,
220
+ value,
221
+ query_chunk_size=query_chunk_size,
222
+ kv_chunk_size=kv_chunk_size,
223
+ kv_chunk_size_min=kv_chunk_size_min,
224
+ use_checkpoint=self.training,
225
+ upcast_attention=upcast_attention,
226
+ )
227
+
228
+ hidden_states = hidden_states.to(dtype)
229
+
230
+ hidden_states = hidden_states.unflatten(0, (-1, self.heads)).transpose(1,2).flatten(start_dim=2)
231
+
232
+ out_proj, dropout = self.to_out
233
+ hidden_states = out_proj(hidden_states)
234
+ hidden_states = dropout(hidden_states)
235
+
236
+ return hidden_states
237
+
238
+
239
+ class CrossAttentionDoggettx(nn.Module):
240
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
241
+ super().__init__()
242
+ inner_dim = dim_head * heads
243
+ context_dim = default(context_dim, query_dim)
244
+
245
+ self.scale = dim_head ** -0.5
246
+ self.heads = heads
247
+
248
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
249
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
250
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
251
+
252
+ self.to_out = nn.Sequential(
253
+ nn.Linear(inner_dim, query_dim),
254
+ nn.Dropout(dropout)
255
+ )
256
+
257
+ def forward(self, x, context=None, mask=None):
258
+ h = self.heads
259
+
260
+ q_in = self.to_q(x)
261
+ context = default(context, x)
262
+ k_in = self.to_k(context)
263
+ v_in = self.to_v(context)
264
+ del context, x
265
+
266
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
267
+ del q_in, k_in, v_in
268
+
269
+ r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
270
+
271
+ mem_free_total = model_management.get_free_memory(q.device)
272
+
273
+ gb = 1024 ** 3
274
+ tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
275
+ modifier = 3 if q.element_size() == 2 else 2.5
276
+ mem_required = tensor_size * modifier
277
+ steps = 1
278
+
279
+
280
+ if mem_required > mem_free_total:
281
+ steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
282
+ # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
283
+ # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
284
+
285
+ if steps > 64:
286
+ max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
287
+ raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
288
+ f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
289
+
290
+ # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size)
291
+ first_op_done = False
292
+ cleared_cache = False
293
+ while True:
294
+ try:
295
+ slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
296
+ for i in range(0, q.shape[1], slice_size):
297
+ end = i + slice_size
298
+ if _ATTN_PRECISION =="fp32":
299
+ with torch.autocast(enabled=False, device_type = 'cuda'):
300
+ s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * self.scale
301
+ else:
302
+ s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
303
+ first_op_done = True
304
+
305
+ s2 = s1.softmax(dim=-1)
306
+ del s1
307
+
308
+ r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
309
+ del s2
310
+ break
311
+ except model_management.OOM_EXCEPTION as e:
312
+ if first_op_done == False:
313
+ torch.cuda.empty_cache()
314
+ torch.cuda.ipc_collect()
315
+ if cleared_cache == False:
316
+ cleared_cache = True
317
+ print("out of memory error, emptying cache and trying again")
318
+ continue
319
+ steps *= 2
320
+ if steps > 64:
321
+ raise e
322
+ print("out of memory error, increasing steps and trying again", steps)
323
+ else:
324
+ raise e
325
+
326
+ del q, k, v
327
+
328
+ r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
329
+ del r1
330
+
331
+ return self.to_out(r2)
332
+
333
+ class CrossAttention(nn.Module):
334
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
335
+ super().__init__()
336
+ inner_dim = dim_head * heads
337
+ context_dim = default(context_dim, query_dim)
338
+
339
+ self.scale = dim_head ** -0.5
340
+ self.heads = heads
341
+
342
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
343
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
344
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
345
+
346
+ self.to_out = nn.Sequential(
347
+ nn.Linear(inner_dim, query_dim),
348
+ nn.Dropout(dropout)
349
+ )
350
+
351
+ def forward(self, x, context=None, mask=None):
352
+ h = self.heads
353
+
354
+ q = self.to_q(x)
355
+ context = default(context, x)
356
+ k = self.to_k(context)
357
+ v = self.to_v(context)
358
+
359
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
360
+
361
+ # force cast to fp32 to avoid overflowing
362
+ if _ATTN_PRECISION =="fp32":
363
+ with torch.autocast(enabled=False, device_type = 'cuda'):
364
+ q, k = q.float(), k.float()
365
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
366
+ else:
367
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
368
+
369
+ del q, k
370
+
371
+ if exists(mask):
372
+ mask = rearrange(mask, 'b ... -> b (...)')
373
+ max_neg_value = -torch.finfo(sim.dtype).max
374
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
375
+ sim.masked_fill_(~mask, max_neg_value)
376
+
377
+ # attention, what we cannot get enough of
378
+ sim = sim.softmax(dim=-1)
379
+
380
+ out = einsum('b i j, b j d -> b i d', sim, v)
381
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
382
+ return self.to_out(out)
383
+
384
+ class MemoryEfficientCrossAttention(nn.Module):
385
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
386
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
387
+ super().__init__()
388
+ print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
389
+ f"{heads} heads.")
390
+ inner_dim = dim_head * heads
391
+ context_dim = default(context_dim, query_dim)
392
+
393
+ self.heads = heads
394
+ self.dim_head = dim_head
395
+
396
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
397
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
398
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
399
+
400
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
401
+ self.attention_op: Optional[Any] = None
402
+
403
+ def forward(self, x, context=None, mask=None):
404
+ q = self.to_q(x)
405
+ context = default(context, x)
406
+ k = self.to_k(context)
407
+ v = self.to_v(context)
408
+
409
+ b, _, _ = q.shape
410
+ q, k, v = map(
411
+ lambda t: t.unsqueeze(3)
412
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
413
+ .permute(0, 2, 1, 3)
414
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
415
+ .contiguous(),
416
+ (q, k, v),
417
+ )
418
+
419
+ # actually compute the attention, what we cannot get enough of
420
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
421
+
422
+ if exists(mask):
423
+ raise NotImplementedError
424
+ out = (
425
+ out.unsqueeze(0)
426
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
427
+ .permute(0, 2, 1, 3)
428
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
429
+ )
430
+ return self.to_out(out)
431
+
432
+ class CrossAttentionPytorch(nn.Module):
433
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
434
+ super().__init__()
435
+ inner_dim = dim_head * heads
436
+ context_dim = default(context_dim, query_dim)
437
+
438
+ self.heads = heads
439
+ self.dim_head = dim_head
440
+
441
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
442
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
443
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
444
+
445
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
446
+ self.attention_op: Optional[Any] = None
447
+
448
+ def forward(self, x, context=None, mask=None):
449
+ q = self.to_q(x)
450
+ context = default(context, x)
451
+ k = self.to_k(context)
452
+ v = self.to_v(context)
453
+
454
+ b, _, _ = q.shape
455
+ q, k, v = map(
456
+ lambda t: t.unsqueeze(3)
457
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
458
+ .permute(0, 2, 1, 3)
459
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
460
+ .contiguous(),
461
+ (q, k, v),
462
+ )
463
+
464
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
465
+
466
+ if exists(mask):
467
+ raise NotImplementedError
468
+ out = (
469
+ out.unsqueeze(0)
470
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
471
+ .permute(0, 2, 1, 3)
472
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
473
+ )
474
+
475
+ return self.to_out(out)
476
+
477
+ import sys
478
+ if model_management.xformers_enabled():
479
+ print("Using xformers cross attention")
480
+ CrossAttention = MemoryEfficientCrossAttention
481
+ elif model_management.pytorch_attention_enabled():
482
+ print("Using pytorch cross attention")
483
+ CrossAttention = CrossAttentionPytorch
484
+ else:
485
+ if "--use-split-cross-attention" in sys.argv:
486
+ print("Using split optimization for cross attention")
487
+ CrossAttention = CrossAttentionDoggettx
488
+ else:
489
+ print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention")
490
+ CrossAttention = CrossAttentionBirchSan
491
+
492
+
493
+ class BasicTransformerBlock(nn.Module):
494
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
495
+ disable_self_attn=False):
496
+ super().__init__()
497
+ self.disable_self_attn = disable_self_attn
498
+ self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
499
+ context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn
500
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
501
+ self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
502
+ heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
503
+ self.norm1 = nn.LayerNorm(dim)
504
+ self.norm2 = nn.LayerNorm(dim)
505
+ self.norm3 = nn.LayerNorm(dim)
506
+ self.checkpoint = checkpoint
507
+
508
+ def forward(self, x, context=None, transformer_options={}):
509
+ return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint)
510
+
511
+ def _forward(self, x, context=None, transformer_options={}):
512
+ n = self.norm1(x)
513
+ if "tomesd" in transformer_options:
514
+ m, u = tomesd.get_functions(x, transformer_options["tomesd"]["ratio"], transformer_options["original_shape"])
515
+ n = u(self.attn1(m(n), context=context if self.disable_self_attn else None))
516
+ else:
517
+ n = self.attn1(n, context=context if self.disable_self_attn else None)
518
+
519
+ x += n
520
+ n = self.norm2(x)
521
+ n = self.attn2(n, context=context)
522
+
523
+ x += n
524
+ x = self.ff(self.norm3(x)) + x
525
+ return x
526
+
527
+
528
+ class SpatialTransformer(nn.Module):
529
+ """
530
+ Transformer block for image-like data.
531
+ First, project the input (aka embedding)
532
+ and reshape to b, t, d.
533
+ Then apply standard transformer action.
534
+ Finally, reshape to image
535
+ NEW: use_linear for more efficiency instead of the 1x1 convs
536
+ """
537
+ def __init__(self, in_channels, n_heads, d_head,
538
+ depth=1, dropout=0., context_dim=None,
539
+ disable_self_attn=False, use_linear=False,
540
+ use_checkpoint=True):
541
+ super().__init__()
542
+ if exists(context_dim) and not isinstance(context_dim, list):
543
+ context_dim = [context_dim]
544
+ self.in_channels = in_channels
545
+ inner_dim = n_heads * d_head
546
+ self.norm = Normalize(in_channels)
547
+ if not use_linear:
548
+ self.proj_in = nn.Conv2d(in_channels,
549
+ inner_dim,
550
+ kernel_size=1,
551
+ stride=1,
552
+ padding=0)
553
+ else:
554
+ self.proj_in = nn.Linear(in_channels, inner_dim)
555
+
556
+ self.transformer_blocks = nn.ModuleList(
557
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
558
+ disable_self_attn=disable_self_attn, checkpoint=use_checkpoint)
559
+ for d in range(depth)]
560
+ )
561
+ if not use_linear:
562
+ self.proj_out = zero_module(nn.Conv2d(inner_dim,
563
+ in_channels,
564
+ kernel_size=1,
565
+ stride=1,
566
+ padding=0))
567
+ else:
568
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
569
+ self.use_linear = use_linear
570
+
571
+ def forward(self, x, context=None, transformer_options={}):
572
+ # note: if no context is given, cross-attention defaults to self-attention
573
+ if not isinstance(context, list):
574
+ context = [context]
575
+ b, c, h, w = x.shape
576
+ x_in = x
577
+ x = self.norm(x)
578
+ if not self.use_linear:
579
+ x = self.proj_in(x)
580
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
581
+ if self.use_linear:
582
+ x = self.proj_in(x)
583
+ for i, block in enumerate(self.transformer_blocks):
584
+ x = block(x, context=context[i], transformer_options=transformer_options)
585
+ if self.use_linear:
586
+ x = self.proj_out(x)
587
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
588
+ if not self.use_linear:
589
+ x = self.proj_out(x)
590
+ return x + x_in
591
+
comfy/ldm/modules/diffusionmodules/__init__.py ADDED
File without changes
comfy/ldm/modules/diffusionmodules/model.py ADDED
@@ -0,0 +1,942 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ import model_management
11
+
12
+ if model_management.xformers_enabled():
13
+ import xformers
14
+ import xformers.ops
15
+
16
+ def get_timestep_embedding(timesteps, embedding_dim):
17
+ """
18
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
19
+ From Fairseq.
20
+ Build sinusoidal embeddings.
21
+ This matches the implementation in tensor2tensor, but differs slightly
22
+ from the description in Section 3.5 of "Attention Is All You Need".
23
+ """
24
+ assert len(timesteps.shape) == 1
25
+
26
+ half_dim = embedding_dim // 2
27
+ emb = math.log(10000) / (half_dim - 1)
28
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
29
+ emb = emb.to(device=timesteps.device)
30
+ emb = timesteps.float()[:, None] * emb[None, :]
31
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
32
+ if embedding_dim % 2 == 1: # zero pad
33
+ emb = torch.nn.functional.pad(emb, (0,1,0,0))
34
+ return emb
35
+
36
+
37
+ def nonlinearity(x):
38
+ # swish
39
+ return x*torch.sigmoid(x)
40
+
41
+
42
+ def Normalize(in_channels, num_groups=32):
43
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
44
+
45
+
46
+ class Upsample(nn.Module):
47
+ def __init__(self, in_channels, with_conv):
48
+ super().__init__()
49
+ self.with_conv = with_conv
50
+ if self.with_conv:
51
+ self.conv = torch.nn.Conv2d(in_channels,
52
+ in_channels,
53
+ kernel_size=3,
54
+ stride=1,
55
+ padding=1)
56
+
57
+ def forward(self, x):
58
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
59
+ if self.with_conv:
60
+ x = self.conv(x)
61
+ return x
62
+
63
+
64
+ class Downsample(nn.Module):
65
+ def __init__(self, in_channels, with_conv):
66
+ super().__init__()
67
+ self.with_conv = with_conv
68
+ if self.with_conv:
69
+ # no asymmetric padding in torch conv, must do it ourselves
70
+ self.conv = torch.nn.Conv2d(in_channels,
71
+ in_channels,
72
+ kernel_size=3,
73
+ stride=2,
74
+ padding=0)
75
+
76
+ def forward(self, x, already_padded=False):
77
+ if self.with_conv:
78
+ if not already_padded:
79
+ pad = (0,1,0,1)
80
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
81
+ x = self.conv(x)
82
+ else:
83
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
84
+ return x
85
+
86
+
87
+ class ResnetBlock(nn.Module):
88
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
89
+ dropout, temb_channels=512):
90
+ super().__init__()
91
+ self.in_channels = in_channels
92
+ out_channels = in_channels if out_channels is None else out_channels
93
+ self.out_channels = out_channels
94
+ self.use_conv_shortcut = conv_shortcut
95
+
96
+ self.swish = torch.nn.SiLU(inplace=True)
97
+ self.norm1 = Normalize(in_channels)
98
+ self.conv1 = torch.nn.Conv2d(in_channels,
99
+ out_channels,
100
+ kernel_size=3,
101
+ stride=1,
102
+ padding=1)
103
+ if temb_channels > 0:
104
+ self.temb_proj = torch.nn.Linear(temb_channels,
105
+ out_channels)
106
+ self.norm2 = Normalize(out_channels)
107
+ self.dropout = torch.nn.Dropout(dropout, inplace=True)
108
+ self.conv2 = torch.nn.Conv2d(out_channels,
109
+ out_channels,
110
+ kernel_size=3,
111
+ stride=1,
112
+ padding=1)
113
+ if self.in_channels != self.out_channels:
114
+ if self.use_conv_shortcut:
115
+ self.conv_shortcut = torch.nn.Conv2d(in_channels,
116
+ out_channels,
117
+ kernel_size=3,
118
+ stride=1,
119
+ padding=1)
120
+ else:
121
+ self.nin_shortcut = torch.nn.Conv2d(in_channels,
122
+ out_channels,
123
+ kernel_size=1,
124
+ stride=1,
125
+ padding=0)
126
+
127
+ def forward(self, x, temb):
128
+ h = x
129
+ h = self.norm1(h)
130
+ h = self.swish(h)
131
+ h = self.conv1(h)
132
+
133
+ if temb is not None:
134
+ h = h + self.temb_proj(self.swish(temb))[:,:,None,None]
135
+
136
+ h = self.norm2(h)
137
+ h = self.swish(h)
138
+ h = self.dropout(h)
139
+ h = self.conv2(h)
140
+
141
+ if self.in_channels != self.out_channels:
142
+ if self.use_conv_shortcut:
143
+ x = self.conv_shortcut(x)
144
+ else:
145
+ x = self.nin_shortcut(x)
146
+
147
+ return x+h
148
+
149
+
150
+ class AttnBlock(nn.Module):
151
+ def __init__(self, in_channels):
152
+ super().__init__()
153
+ self.in_channels = in_channels
154
+
155
+ self.norm = Normalize(in_channels)
156
+ self.q = torch.nn.Conv2d(in_channels,
157
+ in_channels,
158
+ kernel_size=1,
159
+ stride=1,
160
+ padding=0)
161
+ self.k = torch.nn.Conv2d(in_channels,
162
+ in_channels,
163
+ kernel_size=1,
164
+ stride=1,
165
+ padding=0)
166
+ self.v = torch.nn.Conv2d(in_channels,
167
+ in_channels,
168
+ kernel_size=1,
169
+ stride=1,
170
+ padding=0)
171
+ self.proj_out = torch.nn.Conv2d(in_channels,
172
+ in_channels,
173
+ kernel_size=1,
174
+ stride=1,
175
+ padding=0)
176
+
177
+ def forward(self, x):
178
+ h_ = x
179
+ h_ = self.norm(h_)
180
+ q = self.q(h_)
181
+ k = self.k(h_)
182
+ v = self.v(h_)
183
+
184
+ # compute attention
185
+ b,c,h,w = q.shape
186
+ scale = (int(c)**(-0.5))
187
+
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
+ v = v.reshape(b,c,h*w)
192
+
193
+ r1 = torch.zeros_like(k, device=q.device)
194
+
195
+ mem_free_total = model_management.get_free_memory(q.device)
196
+
197
+ gb = 1024 ** 3
198
+ tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
199
+ modifier = 3 if q.element_size() == 2 else 2.5
200
+ mem_required = tensor_size * modifier
201
+ steps = 1
202
+
203
+ if mem_required > mem_free_total:
204
+ steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
205
+
206
+ while True:
207
+ try:
208
+ slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
209
+ for i in range(0, q.shape[1], slice_size):
210
+ end = i + slice_size
211
+ s1 = torch.bmm(q[:, i:end], k) * scale
212
+
213
+ s2 = torch.nn.functional.softmax(s1, dim=2).permute(0,2,1)
214
+ del s1
215
+
216
+ r1[:, :, i:end] = torch.bmm(v, s2)
217
+ del s2
218
+ break
219
+ except model_management.OOM_EXCEPTION as e:
220
+ steps *= 2
221
+ if steps > 128:
222
+ raise e
223
+ print("out of memory error, increasing steps and trying again", steps)
224
+
225
+ h_ = r1.reshape(b,c,h,w)
226
+ del r1
227
+
228
+ h_ = self.proj_out(h_)
229
+
230
+ return x+h_
231
+
232
+ class MemoryEfficientAttnBlock(nn.Module):
233
+ """
234
+ Uses xformers efficient implementation,
235
+ see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
236
+ Note: this is a single-head self-attention operation
237
+ """
238
+ #
239
+ def __init__(self, in_channels):
240
+ super().__init__()
241
+ self.in_channels = in_channels
242
+
243
+ self.norm = Normalize(in_channels)
244
+ self.q = torch.nn.Conv2d(in_channels,
245
+ in_channels,
246
+ kernel_size=1,
247
+ stride=1,
248
+ padding=0)
249
+ self.k = torch.nn.Conv2d(in_channels,
250
+ in_channels,
251
+ kernel_size=1,
252
+ stride=1,
253
+ padding=0)
254
+ self.v = torch.nn.Conv2d(in_channels,
255
+ in_channels,
256
+ kernel_size=1,
257
+ stride=1,
258
+ padding=0)
259
+ self.proj_out = torch.nn.Conv2d(in_channels,
260
+ in_channels,
261
+ kernel_size=1,
262
+ stride=1,
263
+ padding=0)
264
+ self.attention_op: Optional[Any] = None
265
+
266
+ def forward(self, x):
267
+ h_ = x
268
+ h_ = self.norm(h_)
269
+ q = self.q(h_)
270
+ k = self.k(h_)
271
+ v = self.v(h_)
272
+
273
+ # compute attention
274
+ B, C, H, W = q.shape
275
+ q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
276
+
277
+ q, k, v = map(
278
+ lambda t: t.unsqueeze(3)
279
+ .reshape(B, t.shape[1], 1, C)
280
+ .permute(0, 2, 1, 3)
281
+ .reshape(B * 1, t.shape[1], C)
282
+ .contiguous(),
283
+ (q, k, v),
284
+ )
285
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
286
+
287
+ out = (
288
+ out.unsqueeze(0)
289
+ .reshape(B, 1, out.shape[1], C)
290
+ .permute(0, 2, 1, 3)
291
+ .reshape(B, out.shape[1], C)
292
+ )
293
+ out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
294
+ out = self.proj_out(out)
295
+ return x+out
296
+
297
+ class MemoryEfficientAttnBlockPytorch(nn.Module):
298
+ def __init__(self, in_channels):
299
+ super().__init__()
300
+ self.in_channels = in_channels
301
+
302
+ self.norm = Normalize(in_channels)
303
+ self.q = torch.nn.Conv2d(in_channels,
304
+ in_channels,
305
+ kernel_size=1,
306
+ stride=1,
307
+ padding=0)
308
+ self.k = torch.nn.Conv2d(in_channels,
309
+ in_channels,
310
+ kernel_size=1,
311
+ stride=1,
312
+ padding=0)
313
+ self.v = torch.nn.Conv2d(in_channels,
314
+ in_channels,
315
+ kernel_size=1,
316
+ stride=1,
317
+ padding=0)
318
+ self.proj_out = torch.nn.Conv2d(in_channels,
319
+ in_channels,
320
+ kernel_size=1,
321
+ stride=1,
322
+ padding=0)
323
+ self.attention_op: Optional[Any] = None
324
+
325
+ def forward(self, x):
326
+ h_ = x
327
+ h_ = self.norm(h_)
328
+ q = self.q(h_)
329
+ k = self.k(h_)
330
+ v = self.v(h_)
331
+
332
+ # compute attention
333
+ B, C, H, W = q.shape
334
+ q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
335
+
336
+ q, k, v = map(
337
+ lambda t: t.unsqueeze(3)
338
+ .reshape(B, t.shape[1], 1, C)
339
+ .permute(0, 2, 1, 3)
340
+ .reshape(B * 1, t.shape[1], C)
341
+ .contiguous(),
342
+ (q, k, v),
343
+ )
344
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
345
+
346
+ out = (
347
+ out.unsqueeze(0)
348
+ .reshape(B, 1, out.shape[1], C)
349
+ .permute(0, 2, 1, 3)
350
+ .reshape(B, out.shape[1], C)
351
+ )
352
+ out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C)
353
+ out = self.proj_out(out)
354
+ return x+out
355
+
356
+ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
357
+ def forward(self, x, context=None, mask=None):
358
+ b, c, h, w = x.shape
359
+ x = rearrange(x, 'b c h w -> b (h w) c')
360
+ out = super().forward(x, context=context, mask=mask)
361
+ out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w, c=c)
362
+ return x + out
363
+
364
+
365
+ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
366
+ assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown'
367
+ if model_management.xformers_enabled() and attn_type == "vanilla":
368
+ attn_type = "vanilla-xformers"
369
+ if model_management.pytorch_attention_enabled() and attn_type == "vanilla":
370
+ attn_type = "vanilla-pytorch"
371
+ print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
372
+ if attn_type == "vanilla":
373
+ assert attn_kwargs is None
374
+ return AttnBlock(in_channels)
375
+ elif attn_type == "vanilla-xformers":
376
+ print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
377
+ return MemoryEfficientAttnBlock(in_channels)
378
+ elif attn_type == "vanilla-pytorch":
379
+ return MemoryEfficientAttnBlockPytorch(in_channels)
380
+ elif type == "memory-efficient-cross-attn":
381
+ attn_kwargs["query_dim"] = in_channels
382
+ return MemoryEfficientCrossAttentionWrapper(**attn_kwargs)
383
+ elif attn_type == "none":
384
+ return nn.Identity(in_channels)
385
+ else:
386
+ raise NotImplementedError()
387
+
388
+
389
+ class Model(nn.Module):
390
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
391
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
392
+ resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
393
+ super().__init__()
394
+ if use_linear_attn: attn_type = "linear"
395
+ self.ch = ch
396
+ self.temb_ch = self.ch*4
397
+ self.num_resolutions = len(ch_mult)
398
+ self.num_res_blocks = num_res_blocks
399
+ self.resolution = resolution
400
+ self.in_channels = in_channels
401
+
402
+ self.use_timestep = use_timestep
403
+ if self.use_timestep:
404
+ # timestep embedding
405
+ self.temb = nn.Module()
406
+ self.temb.dense = nn.ModuleList([
407
+ torch.nn.Linear(self.ch,
408
+ self.temb_ch),
409
+ torch.nn.Linear(self.temb_ch,
410
+ self.temb_ch),
411
+ ])
412
+
413
+ # downsampling
414
+ self.conv_in = torch.nn.Conv2d(in_channels,
415
+ self.ch,
416
+ kernel_size=3,
417
+ stride=1,
418
+ padding=1)
419
+
420
+ curr_res = resolution
421
+ in_ch_mult = (1,)+tuple(ch_mult)
422
+ self.down = nn.ModuleList()
423
+ for i_level in range(self.num_resolutions):
424
+ block = nn.ModuleList()
425
+ attn = nn.ModuleList()
426
+ block_in = ch*in_ch_mult[i_level]
427
+ block_out = ch*ch_mult[i_level]
428
+ for i_block in range(self.num_res_blocks):
429
+ block.append(ResnetBlock(in_channels=block_in,
430
+ out_channels=block_out,
431
+ temb_channels=self.temb_ch,
432
+ dropout=dropout))
433
+ block_in = block_out
434
+ if curr_res in attn_resolutions:
435
+ attn.append(make_attn(block_in, attn_type=attn_type))
436
+ down = nn.Module()
437
+ down.block = block
438
+ down.attn = attn
439
+ if i_level != self.num_resolutions-1:
440
+ down.downsample = Downsample(block_in, resamp_with_conv)
441
+ curr_res = curr_res // 2
442
+ self.down.append(down)
443
+
444
+ # middle
445
+ self.mid = nn.Module()
446
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
447
+ out_channels=block_in,
448
+ temb_channels=self.temb_ch,
449
+ dropout=dropout)
450
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
451
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
452
+ out_channels=block_in,
453
+ temb_channels=self.temb_ch,
454
+ dropout=dropout)
455
+
456
+ # upsampling
457
+ self.up = nn.ModuleList()
458
+ for i_level in reversed(range(self.num_resolutions)):
459
+ block = nn.ModuleList()
460
+ attn = nn.ModuleList()
461
+ block_out = ch*ch_mult[i_level]
462
+ skip_in = ch*ch_mult[i_level]
463
+ for i_block in range(self.num_res_blocks+1):
464
+ if i_block == self.num_res_blocks:
465
+ skip_in = ch*in_ch_mult[i_level]
466
+ block.append(ResnetBlock(in_channels=block_in+skip_in,
467
+ out_channels=block_out,
468
+ temb_channels=self.temb_ch,
469
+ dropout=dropout))
470
+ block_in = block_out
471
+ if curr_res in attn_resolutions:
472
+ attn.append(make_attn(block_in, attn_type=attn_type))
473
+ up = nn.Module()
474
+ up.block = block
475
+ up.attn = attn
476
+ if i_level != 0:
477
+ up.upsample = Upsample(block_in, resamp_with_conv)
478
+ curr_res = curr_res * 2
479
+ self.up.insert(0, up) # prepend to get consistent order
480
+
481
+ # end
482
+ self.norm_out = Normalize(block_in)
483
+ self.conv_out = torch.nn.Conv2d(block_in,
484
+ out_ch,
485
+ kernel_size=3,
486
+ stride=1,
487
+ padding=1)
488
+
489
+ def forward(self, x, t=None, context=None):
490
+ #assert x.shape[2] == x.shape[3] == self.resolution
491
+ if context is not None:
492
+ # assume aligned context, cat along channel axis
493
+ x = torch.cat((x, context), dim=1)
494
+ if self.use_timestep:
495
+ # timestep embedding
496
+ assert t is not None
497
+ temb = get_timestep_embedding(t, self.ch)
498
+ temb = self.temb.dense[0](temb)
499
+ temb = nonlinearity(temb)
500
+ temb = self.temb.dense[1](temb)
501
+ else:
502
+ temb = None
503
+
504
+ # downsampling
505
+ hs = [self.conv_in(x)]
506
+ for i_level in range(self.num_resolutions):
507
+ for i_block in range(self.num_res_blocks):
508
+ h = self.down[i_level].block[i_block](hs[-1], temb)
509
+ if len(self.down[i_level].attn) > 0:
510
+ h = self.down[i_level].attn[i_block](h)
511
+ hs.append(h)
512
+ if i_level != self.num_resolutions-1:
513
+ hs.append(self.down[i_level].downsample(hs[-1]))
514
+
515
+ # middle
516
+ h = hs[-1]
517
+ h = self.mid.block_1(h, temb)
518
+ h = self.mid.attn_1(h)
519
+ h = self.mid.block_2(h, temb)
520
+
521
+ # upsampling
522
+ for i_level in reversed(range(self.num_resolutions)):
523
+ for i_block in range(self.num_res_blocks+1):
524
+ h = self.up[i_level].block[i_block](
525
+ torch.cat([h, hs.pop()], dim=1), temb)
526
+ if len(self.up[i_level].attn) > 0:
527
+ h = self.up[i_level].attn[i_block](h)
528
+ if i_level != 0:
529
+ h = self.up[i_level].upsample(h)
530
+
531
+ # end
532
+ h = self.norm_out(h)
533
+ h = nonlinearity(h)
534
+ h = self.conv_out(h)
535
+ return h
536
+
537
+ def get_last_layer(self):
538
+ return self.conv_out.weight
539
+
540
+
541
+ class Encoder(nn.Module):
542
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
543
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
544
+ resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
545
+ **ignore_kwargs):
546
+ super().__init__()
547
+ if use_linear_attn: attn_type = "linear"
548
+ self.ch = ch
549
+ self.temb_ch = 0
550
+ self.num_resolutions = len(ch_mult)
551
+ self.num_res_blocks = num_res_blocks
552
+ self.resolution = resolution
553
+ self.in_channels = in_channels
554
+
555
+ # downsampling
556
+ self.conv_in = torch.nn.Conv2d(in_channels,
557
+ self.ch,
558
+ kernel_size=3,
559
+ stride=1,
560
+ padding=1)
561
+
562
+ curr_res = resolution
563
+ in_ch_mult = (1,)+tuple(ch_mult)
564
+ self.in_ch_mult = in_ch_mult
565
+ self.down = nn.ModuleList()
566
+ for i_level in range(self.num_resolutions):
567
+ block = nn.ModuleList()
568
+ attn = nn.ModuleList()
569
+ block_in = ch*in_ch_mult[i_level]
570
+ block_out = ch*ch_mult[i_level]
571
+ for i_block in range(self.num_res_blocks):
572
+ block.append(ResnetBlock(in_channels=block_in,
573
+ out_channels=block_out,
574
+ temb_channels=self.temb_ch,
575
+ dropout=dropout))
576
+ block_in = block_out
577
+ if curr_res in attn_resolutions:
578
+ attn.append(make_attn(block_in, attn_type=attn_type))
579
+ down = nn.Module()
580
+ down.block = block
581
+ down.attn = attn
582
+ if i_level != self.num_resolutions-1:
583
+ down.downsample = Downsample(block_in, resamp_with_conv)
584
+ curr_res = curr_res // 2
585
+ self.down.append(down)
586
+
587
+ # middle
588
+ self.mid = nn.Module()
589
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
590
+ out_channels=block_in,
591
+ temb_channels=self.temb_ch,
592
+ dropout=dropout)
593
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
594
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
595
+ out_channels=block_in,
596
+ temb_channels=self.temb_ch,
597
+ dropout=dropout)
598
+
599
+ # end
600
+ self.norm_out = Normalize(block_in)
601
+ self.conv_out = torch.nn.Conv2d(block_in,
602
+ 2*z_channels if double_z else z_channels,
603
+ kernel_size=3,
604
+ stride=1,
605
+ padding=1)
606
+
607
+ def forward(self, x):
608
+ # timestep embedding
609
+ temb = None
610
+ pad = (0,1,0,1)
611
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
612
+ already_padded = True
613
+ # downsampling
614
+ h = self.conv_in(x)
615
+ for i_level in range(self.num_resolutions):
616
+ for i_block in range(self.num_res_blocks):
617
+ h = self.down[i_level].block[i_block](h, temb)
618
+ if len(self.down[i_level].attn) > 0:
619
+ h = self.down[i_level].attn[i_block](h)
620
+ if i_level != self.num_resolutions-1:
621
+ h = self.down[i_level].downsample(h, already_padded)
622
+ already_padded = False
623
+
624
+ # middle
625
+ h = self.mid.block_1(h, temb)
626
+ h = self.mid.attn_1(h)
627
+ h = self.mid.block_2(h, temb)
628
+
629
+ # end
630
+ h = self.norm_out(h)
631
+ h = nonlinearity(h)
632
+ h = self.conv_out(h)
633
+ return h
634
+
635
+
636
+ class Decoder(nn.Module):
637
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
638
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
639
+ resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
640
+ attn_type="vanilla", **ignorekwargs):
641
+ super().__init__()
642
+ if use_linear_attn: attn_type = "linear"
643
+ self.ch = ch
644
+ self.temb_ch = 0
645
+ self.num_resolutions = len(ch_mult)
646
+ self.num_res_blocks = num_res_blocks
647
+ self.resolution = resolution
648
+ self.in_channels = in_channels
649
+ self.give_pre_end = give_pre_end
650
+ self.tanh_out = tanh_out
651
+
652
+ # compute in_ch_mult, block_in and curr_res at lowest res
653
+ in_ch_mult = (1,)+tuple(ch_mult)
654
+ block_in = ch*ch_mult[self.num_resolutions-1]
655
+ curr_res = resolution // 2**(self.num_resolutions-1)
656
+ self.z_shape = (1,z_channels,curr_res,curr_res)
657
+ print("Working with z of shape {} = {} dimensions.".format(
658
+ self.z_shape, np.prod(self.z_shape)))
659
+
660
+ # z to block_in
661
+ self.conv_in = torch.nn.Conv2d(z_channels,
662
+ block_in,
663
+ kernel_size=3,
664
+ stride=1,
665
+ padding=1)
666
+
667
+ # middle
668
+ self.mid = nn.Module()
669
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
670
+ out_channels=block_in,
671
+ temb_channels=self.temb_ch,
672
+ dropout=dropout)
673
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
674
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
675
+ out_channels=block_in,
676
+ temb_channels=self.temb_ch,
677
+ dropout=dropout)
678
+
679
+ # upsampling
680
+ self.up = nn.ModuleList()
681
+ for i_level in reversed(range(self.num_resolutions)):
682
+ block = nn.ModuleList()
683
+ attn = nn.ModuleList()
684
+ block_out = ch*ch_mult[i_level]
685
+ for i_block in range(self.num_res_blocks+1):
686
+ block.append(ResnetBlock(in_channels=block_in,
687
+ out_channels=block_out,
688
+ temb_channels=self.temb_ch,
689
+ dropout=dropout))
690
+ block_in = block_out
691
+ if curr_res in attn_resolutions:
692
+ attn.append(make_attn(block_in, attn_type=attn_type))
693
+ up = nn.Module()
694
+ up.block = block
695
+ up.attn = attn
696
+ if i_level != 0:
697
+ up.upsample = Upsample(block_in, resamp_with_conv)
698
+ curr_res = curr_res * 2
699
+ self.up.insert(0, up) # prepend to get consistent order
700
+
701
+ # end
702
+ self.norm_out = Normalize(block_in)
703
+ self.conv_out = torch.nn.Conv2d(block_in,
704
+ out_ch,
705
+ kernel_size=3,
706
+ stride=1,
707
+ padding=1)
708
+
709
+ def forward(self, z):
710
+ #assert z.shape[1:] == self.z_shape[1:]
711
+ self.last_z_shape = z.shape
712
+
713
+ # timestep embedding
714
+ temb = None
715
+
716
+ # z to block_in
717
+ h = self.conv_in(z)
718
+
719
+ # middle
720
+ h = self.mid.block_1(h, temb)
721
+ h = self.mid.attn_1(h)
722
+ h = self.mid.block_2(h, temb)
723
+
724
+ # upsampling
725
+ for i_level in reversed(range(self.num_resolutions)):
726
+ for i_block in range(self.num_res_blocks+1):
727
+ h = self.up[i_level].block[i_block](h, temb)
728
+ if len(self.up[i_level].attn) > 0:
729
+ h = self.up[i_level].attn[i_block](h)
730
+ if i_level != 0:
731
+ h = self.up[i_level].upsample(h)
732
+
733
+ # end
734
+ if self.give_pre_end:
735
+ return h
736
+
737
+ h = self.norm_out(h)
738
+ h = nonlinearity(h)
739
+ h = self.conv_out(h)
740
+ if self.tanh_out:
741
+ h = torch.tanh(h)
742
+ return h
743
+
744
+
745
+ class SimpleDecoder(nn.Module):
746
+ def __init__(self, in_channels, out_channels, *args, **kwargs):
747
+ super().__init__()
748
+ self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
749
+ ResnetBlock(in_channels=in_channels,
750
+ out_channels=2 * in_channels,
751
+ temb_channels=0, dropout=0.0),
752
+ ResnetBlock(in_channels=2 * in_channels,
753
+ out_channels=4 * in_channels,
754
+ temb_channels=0, dropout=0.0),
755
+ ResnetBlock(in_channels=4 * in_channels,
756
+ out_channels=2 * in_channels,
757
+ temb_channels=0, dropout=0.0),
758
+ nn.Conv2d(2*in_channels, in_channels, 1),
759
+ Upsample(in_channels, with_conv=True)])
760
+ # end
761
+ self.norm_out = Normalize(in_channels)
762
+ self.conv_out = torch.nn.Conv2d(in_channels,
763
+ out_channels,
764
+ kernel_size=3,
765
+ stride=1,
766
+ padding=1)
767
+
768
+ def forward(self, x):
769
+ for i, layer in enumerate(self.model):
770
+ if i in [1,2,3]:
771
+ x = layer(x, None)
772
+ else:
773
+ x = layer(x)
774
+
775
+ h = self.norm_out(x)
776
+ h = nonlinearity(h)
777
+ x = self.conv_out(h)
778
+ return x
779
+
780
+
781
+ class UpsampleDecoder(nn.Module):
782
+ def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
783
+ ch_mult=(2,2), dropout=0.0):
784
+ super().__init__()
785
+ # upsampling
786
+ self.temb_ch = 0
787
+ self.num_resolutions = len(ch_mult)
788
+ self.num_res_blocks = num_res_blocks
789
+ block_in = in_channels
790
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
791
+ self.res_blocks = nn.ModuleList()
792
+ self.upsample_blocks = nn.ModuleList()
793
+ for i_level in range(self.num_resolutions):
794
+ res_block = []
795
+ block_out = ch * ch_mult[i_level]
796
+ for i_block in range(self.num_res_blocks + 1):
797
+ res_block.append(ResnetBlock(in_channels=block_in,
798
+ out_channels=block_out,
799
+ temb_channels=self.temb_ch,
800
+ dropout=dropout))
801
+ block_in = block_out
802
+ self.res_blocks.append(nn.ModuleList(res_block))
803
+ if i_level != self.num_resolutions - 1:
804
+ self.upsample_blocks.append(Upsample(block_in, True))
805
+ curr_res = curr_res * 2
806
+
807
+ # end
808
+ self.norm_out = Normalize(block_in)
809
+ self.conv_out = torch.nn.Conv2d(block_in,
810
+ out_channels,
811
+ kernel_size=3,
812
+ stride=1,
813
+ padding=1)
814
+
815
+ def forward(self, x):
816
+ # upsampling
817
+ h = x
818
+ for k, i_level in enumerate(range(self.num_resolutions)):
819
+ for i_block in range(self.num_res_blocks + 1):
820
+ h = self.res_blocks[i_level][i_block](h, None)
821
+ if i_level != self.num_resolutions - 1:
822
+ h = self.upsample_blocks[k](h)
823
+ h = self.norm_out(h)
824
+ h = nonlinearity(h)
825
+ h = self.conv_out(h)
826
+ return h
827
+
828
+
829
+ class LatentRescaler(nn.Module):
830
+ def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
831
+ super().__init__()
832
+ # residual block, interpolate, residual block
833
+ self.factor = factor
834
+ self.conv_in = nn.Conv2d(in_channels,
835
+ mid_channels,
836
+ kernel_size=3,
837
+ stride=1,
838
+ padding=1)
839
+ self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
840
+ out_channels=mid_channels,
841
+ temb_channels=0,
842
+ dropout=0.0) for _ in range(depth)])
843
+ self.attn = AttnBlock(mid_channels)
844
+ self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
845
+ out_channels=mid_channels,
846
+ temb_channels=0,
847
+ dropout=0.0) for _ in range(depth)])
848
+
849
+ self.conv_out = nn.Conv2d(mid_channels,
850
+ out_channels,
851
+ kernel_size=1,
852
+ )
853
+
854
+ def forward(self, x):
855
+ x = self.conv_in(x)
856
+ for block in self.res_block1:
857
+ x = block(x, None)
858
+ x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
859
+ x = self.attn(x)
860
+ for block in self.res_block2:
861
+ x = block(x, None)
862
+ x = self.conv_out(x)
863
+ return x
864
+
865
+
866
+ class MergedRescaleEncoder(nn.Module):
867
+ def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
868
+ attn_resolutions, dropout=0.0, resamp_with_conv=True,
869
+ ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
870
+ super().__init__()
871
+ intermediate_chn = ch * ch_mult[-1]
872
+ self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
873
+ z_channels=intermediate_chn, double_z=False, resolution=resolution,
874
+ attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
875
+ out_ch=None)
876
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
877
+ mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
878
+
879
+ def forward(self, x):
880
+ x = self.encoder(x)
881
+ x = self.rescaler(x)
882
+ return x
883
+
884
+
885
+ class MergedRescaleDecoder(nn.Module):
886
+ def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
887
+ dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
888
+ super().__init__()
889
+ tmp_chn = z_channels*ch_mult[-1]
890
+ self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
891
+ resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
892
+ ch_mult=ch_mult, resolution=resolution, ch=ch)
893
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
894
+ out_channels=tmp_chn, depth=rescale_module_depth)
895
+
896
+ def forward(self, x):
897
+ x = self.rescaler(x)
898
+ x = self.decoder(x)
899
+ return x
900
+
901
+
902
+ class Upsampler(nn.Module):
903
+ def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
904
+ super().__init__()
905
+ assert out_size >= in_size
906
+ num_blocks = int(np.log2(out_size//in_size))+1
907
+ factor_up = 1.+ (out_size % in_size)
908
+ print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
909
+ self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
910
+ out_channels=in_channels)
911
+ self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
912
+ attn_resolutions=[], in_channels=None, ch=in_channels,
913
+ ch_mult=[ch_mult for _ in range(num_blocks)])
914
+
915
+ def forward(self, x):
916
+ x = self.rescaler(x)
917
+ x = self.decoder(x)
918
+ return x
919
+
920
+
921
+ class Resize(nn.Module):
922
+ def __init__(self, in_channels=None, learned=False, mode="bilinear"):
923
+ super().__init__()
924
+ self.with_conv = learned
925
+ self.mode = mode
926
+ if self.with_conv:
927
+ print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
928
+ raise NotImplementedError()
929
+ assert in_channels is not None
930
+ # no asymmetric padding in torch conv, must do it ourselves
931
+ self.conv = torch.nn.Conv2d(in_channels,
932
+ in_channels,
933
+ kernel_size=4,
934
+ stride=2,
935
+ padding=1)
936
+
937
+ def forward(self, x, scale_factor=1.0):
938
+ if scale_factor==1.0:
939
+ return x
940
+ else:
941
+ x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
942
+ return x
comfy/ldm/modules/diffusionmodules/openaimodel.py ADDED
@@ -0,0 +1,821 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import abstractmethod
2
+ import math
3
+
4
+ import numpy as np
5
+ import torch as th
6
+ import torch.nn as nn
7
+ import torch.nn.functional as F
8
+
9
+ from ldm.modules.diffusionmodules.util import (
10
+ checkpoint,
11
+ conv_nd,
12
+ linear,
13
+ avg_pool_nd,
14
+ zero_module,
15
+ normalization,
16
+ timestep_embedding,
17
+ )
18
+ from ldm.modules.attention import SpatialTransformer
19
+ from ldm.util import exists
20
+
21
+
22
+ # dummy replace
23
+ def convert_module_to_f16(x):
24
+ pass
25
+
26
+ def convert_module_to_f32(x):
27
+ pass
28
+
29
+
30
+ ## go
31
+ class AttentionPool2d(nn.Module):
32
+ """
33
+ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ spacial_dim: int,
39
+ embed_dim: int,
40
+ num_heads_channels: int,
41
+ output_dim: int = None,
42
+ ):
43
+ super().__init__()
44
+ self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
45
+ self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
46
+ self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
47
+ self.num_heads = embed_dim // num_heads_channels
48
+ self.attention = QKVAttention(self.num_heads)
49
+
50
+ def forward(self, x):
51
+ b, c, *_spatial = x.shape
52
+ x = x.reshape(b, c, -1) # NC(HW)
53
+ x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
54
+ x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
55
+ x = self.qkv_proj(x)
56
+ x = self.attention(x)
57
+ x = self.c_proj(x)
58
+ return x[:, :, 0]
59
+
60
+
61
+ class TimestepBlock(nn.Module):
62
+ """
63
+ Any module where forward() takes timestep embeddings as a second argument.
64
+ """
65
+
66
+ @abstractmethod
67
+ def forward(self, x, emb):
68
+ """
69
+ Apply the module to `x` given `emb` timestep embeddings.
70
+ """
71
+
72
+
73
+ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
74
+ """
75
+ A sequential module that passes timestep embeddings to the children that
76
+ support it as an extra input.
77
+ """
78
+
79
+ def forward(self, x, emb, context=None, transformer_options={}):
80
+ for layer in self:
81
+ if isinstance(layer, TimestepBlock):
82
+ x = layer(x, emb)
83
+ elif isinstance(layer, SpatialTransformer):
84
+ x = layer(x, context, transformer_options)
85
+ else:
86
+ x = layer(x)
87
+ return x
88
+
89
+
90
+ class Upsample(nn.Module):
91
+ """
92
+ An upsampling layer with an optional convolution.
93
+ :param channels: channels in the inputs and outputs.
94
+ :param use_conv: a bool determining if a convolution is applied.
95
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
96
+ upsampling occurs in the inner-two dimensions.
97
+ """
98
+
99
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
100
+ super().__init__()
101
+ self.channels = channels
102
+ self.out_channels = out_channels or channels
103
+ self.use_conv = use_conv
104
+ self.dims = dims
105
+ if use_conv:
106
+ self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
107
+
108
+ def forward(self, x):
109
+ assert x.shape[1] == self.channels
110
+ if self.dims == 3:
111
+ x = F.interpolate(
112
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
113
+ )
114
+ else:
115
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
116
+ if self.use_conv:
117
+ x = self.conv(x)
118
+ return x
119
+
120
+ class TransposedUpsample(nn.Module):
121
+ 'Learned 2x upsampling without padding'
122
+ def __init__(self, channels, out_channels=None, ks=5):
123
+ super().__init__()
124
+ self.channels = channels
125
+ self.out_channels = out_channels or channels
126
+
127
+ self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
128
+
129
+ def forward(self,x):
130
+ return self.up(x)
131
+
132
+
133
+ class Downsample(nn.Module):
134
+ """
135
+ A downsampling layer with an optional convolution.
136
+ :param channels: channels in the inputs and outputs.
137
+ :param use_conv: a bool determining if a convolution is applied.
138
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
139
+ downsampling occurs in the inner-two dimensions.
140
+ """
141
+
142
+ def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1):
143
+ super().__init__()
144
+ self.channels = channels
145
+ self.out_channels = out_channels or channels
146
+ self.use_conv = use_conv
147
+ self.dims = dims
148
+ stride = 2 if dims != 3 else (1, 2, 2)
149
+ if use_conv:
150
+ self.op = conv_nd(
151
+ dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
152
+ )
153
+ else:
154
+ assert self.channels == self.out_channels
155
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
156
+
157
+ def forward(self, x):
158
+ assert x.shape[1] == self.channels
159
+ return self.op(x)
160
+
161
+
162
+ class ResBlock(TimestepBlock):
163
+ """
164
+ A residual block that can optionally change the number of channels.
165
+ :param channels: the number of input channels.
166
+ :param emb_channels: the number of timestep embedding channels.
167
+ :param dropout: the rate of dropout.
168
+ :param out_channels: if specified, the number of out channels.
169
+ :param use_conv: if True and out_channels is specified, use a spatial
170
+ convolution instead of a smaller 1x1 convolution to change the
171
+ channels in the skip connection.
172
+ :param dims: determines if the signal is 1D, 2D, or 3D.
173
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
174
+ :param up: if True, use this block for upsampling.
175
+ :param down: if True, use this block for downsampling.
176
+ """
177
+
178
+ def __init__(
179
+ self,
180
+ channels,
181
+ emb_channels,
182
+ dropout,
183
+ out_channels=None,
184
+ use_conv=False,
185
+ use_scale_shift_norm=False,
186
+ dims=2,
187
+ use_checkpoint=False,
188
+ up=False,
189
+ down=False,
190
+ ):
191
+ super().__init__()
192
+ self.channels = channels
193
+ self.emb_channels = emb_channels
194
+ self.dropout = dropout
195
+ self.out_channels = out_channels or channels
196
+ self.use_conv = use_conv
197
+ self.use_checkpoint = use_checkpoint
198
+ self.use_scale_shift_norm = use_scale_shift_norm
199
+
200
+ self.in_layers = nn.Sequential(
201
+ normalization(channels),
202
+ nn.SiLU(),
203
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
204
+ )
205
+
206
+ self.updown = up or down
207
+
208
+ if up:
209
+ self.h_upd = Upsample(channels, False, dims)
210
+ self.x_upd = Upsample(channels, False, dims)
211
+ elif down:
212
+ self.h_upd = Downsample(channels, False, dims)
213
+ self.x_upd = Downsample(channels, False, dims)
214
+ else:
215
+ self.h_upd = self.x_upd = nn.Identity()
216
+
217
+ self.emb_layers = nn.Sequential(
218
+ nn.SiLU(),
219
+ linear(
220
+ emb_channels,
221
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
222
+ ),
223
+ )
224
+ self.out_layers = nn.Sequential(
225
+ normalization(self.out_channels),
226
+ nn.SiLU(),
227
+ nn.Dropout(p=dropout),
228
+ zero_module(
229
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
230
+ ),
231
+ )
232
+
233
+ if self.out_channels == channels:
234
+ self.skip_connection = nn.Identity()
235
+ elif use_conv:
236
+ self.skip_connection = conv_nd(
237
+ dims, channels, self.out_channels, 3, padding=1
238
+ )
239
+ else:
240
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
241
+
242
+ def forward(self, x, emb):
243
+ """
244
+ Apply the block to a Tensor, conditioned on a timestep embedding.
245
+ :param x: an [N x C x ...] Tensor of features.
246
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
247
+ :return: an [N x C x ...] Tensor of outputs.
248
+ """
249
+ return checkpoint(
250
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
251
+ )
252
+
253
+
254
+ def _forward(self, x, emb):
255
+ if self.updown:
256
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
257
+ h = in_rest(x)
258
+ h = self.h_upd(h)
259
+ x = self.x_upd(x)
260
+ h = in_conv(h)
261
+ else:
262
+ h = self.in_layers(x)
263
+ emb_out = self.emb_layers(emb).type(h.dtype)
264
+ while len(emb_out.shape) < len(h.shape):
265
+ emb_out = emb_out[..., None]
266
+ if self.use_scale_shift_norm:
267
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
268
+ scale, shift = th.chunk(emb_out, 2, dim=1)
269
+ h = out_norm(h) * (1 + scale) + shift
270
+ h = out_rest(h)
271
+ else:
272
+ h = h + emb_out
273
+ h = self.out_layers(h)
274
+ return self.skip_connection(x) + h
275
+
276
+
277
+ class AttentionBlock(nn.Module):
278
+ """
279
+ An attention block that allows spatial positions to attend to each other.
280
+ Originally ported from here, but adapted to the N-d case.
281
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
282
+ """
283
+
284
+ def __init__(
285
+ self,
286
+ channels,
287
+ num_heads=1,
288
+ num_head_channels=-1,
289
+ use_checkpoint=False,
290
+ use_new_attention_order=False,
291
+ ):
292
+ super().__init__()
293
+ self.channels = channels
294
+ if num_head_channels == -1:
295
+ self.num_heads = num_heads
296
+ else:
297
+ assert (
298
+ channels % num_head_channels == 0
299
+ ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
300
+ self.num_heads = channels // num_head_channels
301
+ self.use_checkpoint = use_checkpoint
302
+ self.norm = normalization(channels)
303
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
304
+ if use_new_attention_order:
305
+ # split qkv before split heads
306
+ self.attention = QKVAttention(self.num_heads)
307
+ else:
308
+ # split heads before split qkv
309
+ self.attention = QKVAttentionLegacy(self.num_heads)
310
+
311
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
312
+
313
+ def forward(self, x):
314
+ return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
315
+ #return pt_checkpoint(self._forward, x) # pytorch
316
+
317
+ def _forward(self, x):
318
+ b, c, *spatial = x.shape
319
+ x = x.reshape(b, c, -1)
320
+ qkv = self.qkv(self.norm(x))
321
+ h = self.attention(qkv)
322
+ h = self.proj_out(h)
323
+ return (x + h).reshape(b, c, *spatial)
324
+
325
+
326
+ def count_flops_attn(model, _x, y):
327
+ """
328
+ A counter for the `thop` package to count the operations in an
329
+ attention operation.
330
+ Meant to be used like:
331
+ macs, params = thop.profile(
332
+ model,
333
+ inputs=(inputs, timestamps),
334
+ custom_ops={QKVAttention: QKVAttention.count_flops},
335
+ )
336
+ """
337
+ b, c, *spatial = y[0].shape
338
+ num_spatial = int(np.prod(spatial))
339
+ # We perform two matmuls with the same number of ops.
340
+ # The first computes the weight matrix, the second computes
341
+ # the combination of the value vectors.
342
+ matmul_ops = 2 * b * (num_spatial ** 2) * c
343
+ model.total_ops += th.DoubleTensor([matmul_ops])
344
+
345
+
346
+ class QKVAttentionLegacy(nn.Module):
347
+ """
348
+ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
349
+ """
350
+
351
+ def __init__(self, n_heads):
352
+ super().__init__()
353
+ self.n_heads = n_heads
354
+
355
+ def forward(self, qkv):
356
+ """
357
+ Apply QKV attention.
358
+ :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
359
+ :return: an [N x (H * C) x T] tensor after attention.
360
+ """
361
+ bs, width, length = qkv.shape
362
+ assert width % (3 * self.n_heads) == 0
363
+ ch = width // (3 * self.n_heads)
364
+ q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
365
+ scale = 1 / math.sqrt(math.sqrt(ch))
366
+ weight = th.einsum(
367
+ "bct,bcs->bts", q * scale, k * scale
368
+ ) # More stable with f16 than dividing afterwards
369
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
370
+ a = th.einsum("bts,bcs->bct", weight, v)
371
+ return a.reshape(bs, -1, length)
372
+
373
+ @staticmethod
374
+ def count_flops(model, _x, y):
375
+ return count_flops_attn(model, _x, y)
376
+
377
+
378
+ class QKVAttention(nn.Module):
379
+ """
380
+ A module which performs QKV attention and splits in a different order.
381
+ """
382
+
383
+ def __init__(self, n_heads):
384
+ super().__init__()
385
+ self.n_heads = n_heads
386
+
387
+ def forward(self, qkv):
388
+ """
389
+ Apply QKV attention.
390
+ :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
391
+ :return: an [N x (H * C) x T] tensor after attention.
392
+ """
393
+ bs, width, length = qkv.shape
394
+ assert width % (3 * self.n_heads) == 0
395
+ ch = width // (3 * self.n_heads)
396
+ q, k, v = qkv.chunk(3, dim=1)
397
+ scale = 1 / math.sqrt(math.sqrt(ch))
398
+ weight = th.einsum(
399
+ "bct,bcs->bts",
400
+ (q * scale).view(bs * self.n_heads, ch, length),
401
+ (k * scale).view(bs * self.n_heads, ch, length),
402
+ ) # More stable with f16 than dividing afterwards
403
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
404
+ a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
405
+ return a.reshape(bs, -1, length)
406
+
407
+ @staticmethod
408
+ def count_flops(model, _x, y):
409
+ return count_flops_attn(model, _x, y)
410
+
411
+
412
+ class Timestep(nn.Module):
413
+ def __init__(self, dim):
414
+ super().__init__()
415
+ self.dim = dim
416
+
417
+ def forward(self, t):
418
+ return timestep_embedding(t, self.dim)
419
+
420
+
421
+ class UNetModel(nn.Module):
422
+ """
423
+ The full UNet model with attention and timestep embedding.
424
+ :param in_channels: channels in the input Tensor.
425
+ :param model_channels: base channel count for the model.
426
+ :param out_channels: channels in the output Tensor.
427
+ :param num_res_blocks: number of residual blocks per downsample.
428
+ :param attention_resolutions: a collection of downsample rates at which
429
+ attention will take place. May be a set, list, or tuple.
430
+ For example, if this contains 4, then at 4x downsampling, attention
431
+ will be used.
432
+ :param dropout: the dropout probability.
433
+ :param channel_mult: channel multiplier for each level of the UNet.
434
+ :param conv_resample: if True, use learned convolutions for upsampling and
435
+ downsampling.
436
+ :param dims: determines if the signal is 1D, 2D, or 3D.
437
+ :param num_classes: if specified (as an int), then this model will be
438
+ class-conditional with `num_classes` classes.
439
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
440
+ :param num_heads: the number of attention heads in each attention layer.
441
+ :param num_heads_channels: if specified, ignore num_heads and instead use
442
+ a fixed channel width per attention head.
443
+ :param num_heads_upsample: works with num_heads to set a different number
444
+ of heads for upsampling. Deprecated.
445
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
446
+ :param resblock_updown: use residual blocks for up/downsampling.
447
+ :param use_new_attention_order: use a different attention pattern for potentially
448
+ increased efficiency.
449
+ """
450
+
451
+ def __init__(
452
+ self,
453
+ image_size,
454
+ in_channels,
455
+ model_channels,
456
+ out_channels,
457
+ num_res_blocks,
458
+ attention_resolutions,
459
+ dropout=0,
460
+ channel_mult=(1, 2, 4, 8),
461
+ conv_resample=True,
462
+ dims=2,
463
+ num_classes=None,
464
+ use_checkpoint=False,
465
+ use_fp16=False,
466
+ use_bf16=False,
467
+ num_heads=-1,
468
+ num_head_channels=-1,
469
+ num_heads_upsample=-1,
470
+ use_scale_shift_norm=False,
471
+ resblock_updown=False,
472
+ use_new_attention_order=False,
473
+ use_spatial_transformer=False, # custom transformer support
474
+ transformer_depth=1, # custom transformer support
475
+ context_dim=None, # custom transformer support
476
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
477
+ legacy=True,
478
+ disable_self_attentions=None,
479
+ num_attention_blocks=None,
480
+ disable_middle_self_attn=False,
481
+ use_linear_in_transformer=False,
482
+ adm_in_channels=None,
483
+ ):
484
+ super().__init__()
485
+ if use_spatial_transformer:
486
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
487
+
488
+ if context_dim is not None:
489
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
490
+ # from omegaconf.listconfig import ListConfig
491
+ # if type(context_dim) == ListConfig:
492
+ # context_dim = list(context_dim)
493
+
494
+ if num_heads_upsample == -1:
495
+ num_heads_upsample = num_heads
496
+
497
+ if num_heads == -1:
498
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
499
+
500
+ if num_head_channels == -1:
501
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
502
+
503
+ self.image_size = image_size
504
+ self.in_channels = in_channels
505
+ self.model_channels = model_channels
506
+ self.out_channels = out_channels
507
+ if isinstance(num_res_blocks, int):
508
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
509
+ else:
510
+ if len(num_res_blocks) != len(channel_mult):
511
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
512
+ "as a list/tuple (per-level) with the same length as channel_mult")
513
+ self.num_res_blocks = num_res_blocks
514
+ if disable_self_attentions is not None:
515
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
516
+ assert len(disable_self_attentions) == len(channel_mult)
517
+ if num_attention_blocks is not None:
518
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
519
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
520
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
521
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
522
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
523
+ f"attention will still not be set.")
524
+
525
+ self.attention_resolutions = attention_resolutions
526
+ self.dropout = dropout
527
+ self.channel_mult = channel_mult
528
+ self.conv_resample = conv_resample
529
+ self.num_classes = num_classes
530
+ self.use_checkpoint = use_checkpoint
531
+ self.dtype = th.float16 if use_fp16 else th.float32
532
+ self.dtype = th.bfloat16 if use_bf16 else self.dtype
533
+ self.num_heads = num_heads
534
+ self.num_head_channels = num_head_channels
535
+ self.num_heads_upsample = num_heads_upsample
536
+ self.predict_codebook_ids = n_embed is not None
537
+
538
+ time_embed_dim = model_channels * 4
539
+ self.time_embed = nn.Sequential(
540
+ linear(model_channels, time_embed_dim),
541
+ nn.SiLU(),
542
+ linear(time_embed_dim, time_embed_dim),
543
+ )
544
+
545
+ if self.num_classes is not None:
546
+ if isinstance(self.num_classes, int):
547
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
548
+ elif self.num_classes == "continuous":
549
+ print("setting up linear c_adm embedding layer")
550
+ self.label_emb = nn.Linear(1, time_embed_dim)
551
+ elif self.num_classes == "sequential":
552
+ assert adm_in_channels is not None
553
+ self.label_emb = nn.Sequential(
554
+ nn.Sequential(
555
+ linear(adm_in_channels, time_embed_dim),
556
+ nn.SiLU(),
557
+ linear(time_embed_dim, time_embed_dim),
558
+ )
559
+ )
560
+ else:
561
+ raise ValueError()
562
+
563
+ self.input_blocks = nn.ModuleList(
564
+ [
565
+ TimestepEmbedSequential(
566
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
567
+ )
568
+ ]
569
+ )
570
+ self._feature_size = model_channels
571
+ input_block_chans = [model_channels]
572
+ ch = model_channels
573
+ ds = 1
574
+ for level, mult in enumerate(channel_mult):
575
+ for nr in range(self.num_res_blocks[level]):
576
+ layers = [
577
+ ResBlock(
578
+ ch,
579
+ time_embed_dim,
580
+ dropout,
581
+ out_channels=mult * model_channels,
582
+ dims=dims,
583
+ use_checkpoint=use_checkpoint,
584
+ use_scale_shift_norm=use_scale_shift_norm,
585
+ )
586
+ ]
587
+ ch = mult * model_channels
588
+ if ds in attention_resolutions:
589
+ if num_head_channels == -1:
590
+ dim_head = ch // num_heads
591
+ else:
592
+ num_heads = ch // num_head_channels
593
+ dim_head = num_head_channels
594
+ if legacy:
595
+ #num_heads = 1
596
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
597
+ if exists(disable_self_attentions):
598
+ disabled_sa = disable_self_attentions[level]
599
+ else:
600
+ disabled_sa = False
601
+
602
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
603
+ layers.append(
604
+ AttentionBlock(
605
+ ch,
606
+ use_checkpoint=use_checkpoint,
607
+ num_heads=num_heads,
608
+ num_head_channels=dim_head,
609
+ use_new_attention_order=use_new_attention_order,
610
+ ) if not use_spatial_transformer else SpatialTransformer(
611
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
612
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
613
+ use_checkpoint=use_checkpoint
614
+ )
615
+ )
616
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
617
+ self._feature_size += ch
618
+ input_block_chans.append(ch)
619
+ if level != len(channel_mult) - 1:
620
+ out_ch = ch
621
+ self.input_blocks.append(
622
+ TimestepEmbedSequential(
623
+ ResBlock(
624
+ ch,
625
+ time_embed_dim,
626
+ dropout,
627
+ out_channels=out_ch,
628
+ dims=dims,
629
+ use_checkpoint=use_checkpoint,
630
+ use_scale_shift_norm=use_scale_shift_norm,
631
+ down=True,
632
+ )
633
+ if resblock_updown
634
+ else Downsample(
635
+ ch, conv_resample, dims=dims, out_channels=out_ch
636
+ )
637
+ )
638
+ )
639
+ ch = out_ch
640
+ input_block_chans.append(ch)
641
+ ds *= 2
642
+ self._feature_size += ch
643
+
644
+ if num_head_channels == -1:
645
+ dim_head = ch // num_heads
646
+ else:
647
+ num_heads = ch // num_head_channels
648
+ dim_head = num_head_channels
649
+ if legacy:
650
+ #num_heads = 1
651
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
652
+ self.middle_block = TimestepEmbedSequential(
653
+ ResBlock(
654
+ ch,
655
+ time_embed_dim,
656
+ dropout,
657
+ dims=dims,
658
+ use_checkpoint=use_checkpoint,
659
+ use_scale_shift_norm=use_scale_shift_norm,
660
+ ),
661
+ AttentionBlock(
662
+ ch,
663
+ use_checkpoint=use_checkpoint,
664
+ num_heads=num_heads,
665
+ num_head_channels=dim_head,
666
+ use_new_attention_order=use_new_attention_order,
667
+ ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
668
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
669
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
670
+ use_checkpoint=use_checkpoint
671
+ ),
672
+ ResBlock(
673
+ ch,
674
+ time_embed_dim,
675
+ dropout,
676
+ dims=dims,
677
+ use_checkpoint=use_checkpoint,
678
+ use_scale_shift_norm=use_scale_shift_norm,
679
+ ),
680
+ )
681
+ self._feature_size += ch
682
+
683
+ self.output_blocks = nn.ModuleList([])
684
+ for level, mult in list(enumerate(channel_mult))[::-1]:
685
+ for i in range(self.num_res_blocks[level] + 1):
686
+ ich = input_block_chans.pop()
687
+ layers = [
688
+ ResBlock(
689
+ ch + ich,
690
+ time_embed_dim,
691
+ dropout,
692
+ out_channels=model_channels * mult,
693
+ dims=dims,
694
+ use_checkpoint=use_checkpoint,
695
+ use_scale_shift_norm=use_scale_shift_norm,
696
+ )
697
+ ]
698
+ ch = model_channels * mult
699
+ if ds in attention_resolutions:
700
+ if num_head_channels == -1:
701
+ dim_head = ch // num_heads
702
+ else:
703
+ num_heads = ch // num_head_channels
704
+ dim_head = num_head_channels
705
+ if legacy:
706
+ #num_heads = 1
707
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
708
+ if exists(disable_self_attentions):
709
+ disabled_sa = disable_self_attentions[level]
710
+ else:
711
+ disabled_sa = False
712
+
713
+ if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
714
+ layers.append(
715
+ AttentionBlock(
716
+ ch,
717
+ use_checkpoint=use_checkpoint,
718
+ num_heads=num_heads_upsample,
719
+ num_head_channels=dim_head,
720
+ use_new_attention_order=use_new_attention_order,
721
+ ) if not use_spatial_transformer else SpatialTransformer(
722
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
723
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
724
+ use_checkpoint=use_checkpoint
725
+ )
726
+ )
727
+ if level and i == self.num_res_blocks[level]:
728
+ out_ch = ch
729
+ layers.append(
730
+ ResBlock(
731
+ ch,
732
+ time_embed_dim,
733
+ dropout,
734
+ out_channels=out_ch,
735
+ dims=dims,
736
+ use_checkpoint=use_checkpoint,
737
+ use_scale_shift_norm=use_scale_shift_norm,
738
+ up=True,
739
+ )
740
+ if resblock_updown
741
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
742
+ )
743
+ ds //= 2
744
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
745
+ self._feature_size += ch
746
+
747
+ self.out = nn.Sequential(
748
+ normalization(ch),
749
+ nn.SiLU(),
750
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
751
+ )
752
+ if self.predict_codebook_ids:
753
+ self.id_predictor = nn.Sequential(
754
+ normalization(ch),
755
+ conv_nd(dims, model_channels, n_embed, 1),
756
+ #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
757
+ )
758
+
759
+ def convert_to_fp16(self):
760
+ """
761
+ Convert the torso of the model to float16.
762
+ """
763
+ self.input_blocks.apply(convert_module_to_f16)
764
+ self.middle_block.apply(convert_module_to_f16)
765
+ self.output_blocks.apply(convert_module_to_f16)
766
+
767
+ def convert_to_fp32(self):
768
+ """
769
+ Convert the torso of the model to float32.
770
+ """
771
+ self.input_blocks.apply(convert_module_to_f32)
772
+ self.middle_block.apply(convert_module_to_f32)
773
+ self.output_blocks.apply(convert_module_to_f32)
774
+
775
+ def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
776
+ """
777
+ Apply the model to an input batch.
778
+ :param x: an [N x C x ...] Tensor of inputs.
779
+ :param timesteps: a 1-D batch of timesteps.
780
+ :param context: conditioning plugged in via crossattn
781
+ :param y: an [N] Tensor of labels, if class-conditional.
782
+ :return: an [N x C x ...] Tensor of outputs.
783
+ """
784
+ transformer_options["original_shape"] = list(x.shape)
785
+ assert (y is not None) == (
786
+ self.num_classes is not None
787
+ ), "must specify y if and only if the model is class-conditional"
788
+ hs = []
789
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
790
+ emb = self.time_embed(t_emb)
791
+
792
+ if self.num_classes is not None:
793
+ assert y.shape[0] == x.shape[0]
794
+ emb = emb + self.label_emb(y)
795
+
796
+ h = x.type(self.dtype)
797
+ for id, module in enumerate(self.input_blocks):
798
+ h = module(h, emb, context, transformer_options)
799
+ if control is not None and 'input' in control and len(control['input']) > 0:
800
+ ctrl = control['input'].pop()
801
+ if ctrl is not None:
802
+ h += ctrl
803
+ hs.append(h)
804
+ h = self.middle_block(h, emb, context, transformer_options)
805
+ if control is not None and 'middle' in control and len(control['middle']) > 0:
806
+ h += control['middle'].pop()
807
+
808
+ for module in self.output_blocks:
809
+ hsp = hs.pop()
810
+ if control is not None and 'output' in control and len(control['output']) > 0:
811
+ ctrl = control['output'].pop()
812
+ if ctrl is not None:
813
+ hsp += ctrl
814
+ h = th.cat([h, hsp], dim=1)
815
+ del hsp
816
+ h = module(h, emb, context, transformer_options)
817
+ h = h.type(x.dtype)
818
+ if self.predict_codebook_ids:
819
+ return self.id_predictor(h)
820
+ else:
821
+ return self.out(h)
comfy/ldm/modules/diffusionmodules/upscaling.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import numpy as np
4
+ from functools import partial
5
+
6
+ from ldm.modules.diffusionmodules.util import extract_into_tensor, make_beta_schedule
7
+ from ldm.util import default
8
+
9
+
10
+ class AbstractLowScaleModel(nn.Module):
11
+ # for concatenating a downsampled image to the latent representation
12
+ def __init__(self, noise_schedule_config=None):
13
+ super(AbstractLowScaleModel, self).__init__()
14
+ if noise_schedule_config is not None:
15
+ self.register_schedule(**noise_schedule_config)
16
+
17
+ def register_schedule(self, beta_schedule="linear", timesteps=1000,
18
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
19
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
20
+ cosine_s=cosine_s)
21
+ alphas = 1. - betas
22
+ alphas_cumprod = np.cumprod(alphas, axis=0)
23
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
24
+
25
+ timesteps, = betas.shape
26
+ self.num_timesteps = int(timesteps)
27
+ self.linear_start = linear_start
28
+ self.linear_end = linear_end
29
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
30
+
31
+ to_torch = partial(torch.tensor, dtype=torch.float32)
32
+
33
+ self.register_buffer('betas', to_torch(betas))
34
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
35
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
36
+
37
+ # calculations for diffusion q(x_t | x_{t-1}) and others
38
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
39
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
40
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
41
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
42
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
43
+
44
+ def q_sample(self, x_start, t, noise=None):
45
+ noise = default(noise, lambda: torch.randn_like(x_start))
46
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
47
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
48
+
49
+ def forward(self, x):
50
+ return x, None
51
+
52
+ def decode(self, x):
53
+ return x
54
+
55
+
56
+ class SimpleImageConcat(AbstractLowScaleModel):
57
+ # no noise level conditioning
58
+ def __init__(self):
59
+ super(SimpleImageConcat, self).__init__(noise_schedule_config=None)
60
+ self.max_noise_level = 0
61
+
62
+ def forward(self, x):
63
+ # fix to constant noise level
64
+ return x, torch.zeros(x.shape[0], device=x.device).long()
65
+
66
+
67
+ class ImageConcatWithNoiseAugmentation(AbstractLowScaleModel):
68
+ def __init__(self, noise_schedule_config, max_noise_level=1000, to_cuda=False):
69
+ super().__init__(noise_schedule_config=noise_schedule_config)
70
+ self.max_noise_level = max_noise_level
71
+
72
+ def forward(self, x, noise_level=None):
73
+ if noise_level is None:
74
+ noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
75
+ else:
76
+ assert isinstance(noise_level, torch.Tensor)
77
+ z = self.q_sample(x, noise_level)
78
+ return z, noise_level
79
+
80
+
81
+
comfy/ldm/modules/diffusionmodules/util.py ADDED
@@ -0,0 +1,278 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 == "squaredcos_cap_v2": # used for karlo prior
38
+ # return early
39
+ return betas_for_alpha_bar(
40
+ n_timestep,
41
+ lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
42
+ )
43
+
44
+ elif schedule == "sqrt_linear":
45
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
46
+ elif schedule == "sqrt":
47
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
48
+ else:
49
+ raise ValueError(f"schedule '{schedule}' unknown.")
50
+ return betas.numpy()
51
+
52
+
53
+ def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
54
+ if ddim_discr_method == 'uniform':
55
+ c = num_ddpm_timesteps // num_ddim_timesteps
56
+ ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
57
+ elif ddim_discr_method == 'quad':
58
+ ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
59
+ else:
60
+ raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
61
+
62
+ # assert ddim_timesteps.shape[0] == num_ddim_timesteps
63
+ # add one to get the final alpha values right (the ones from first scale to data during sampling)
64
+ steps_out = ddim_timesteps + 1
65
+ if verbose:
66
+ print(f'Selected timesteps for ddim sampler: {steps_out}')
67
+ return steps_out
68
+
69
+
70
+ def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
71
+ # select alphas for computing the variance schedule
72
+ alphas = alphacums[ddim_timesteps]
73
+ alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
74
+
75
+ # according the the formula provided in https://arxiv.org/abs/2010.02502
76
+ sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
77
+ if verbose:
78
+ print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
79
+ print(f'For the chosen value of eta, which is {eta}, '
80
+ f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
81
+ return sigmas, alphas, alphas_prev
82
+
83
+
84
+ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
85
+ """
86
+ Create a beta schedule that discretizes the given alpha_t_bar function,
87
+ which defines the cumulative product of (1-beta) over time from t = [0,1].
88
+ :param num_diffusion_timesteps: the number of betas to produce.
89
+ :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
90
+ produces the cumulative product of (1-beta) up to that
91
+ part of the diffusion process.
92
+ :param max_beta: the maximum beta to use; use values lower than 1 to
93
+ prevent singularities.
94
+ """
95
+ betas = []
96
+ for i in range(num_diffusion_timesteps):
97
+ t1 = i / num_diffusion_timesteps
98
+ t2 = (i + 1) / num_diffusion_timesteps
99
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
100
+ return np.array(betas)
101
+
102
+
103
+ def extract_into_tensor(a, t, x_shape):
104
+ b, *_ = t.shape
105
+ out = a.gather(-1, t)
106
+ return out.reshape(b, *((1,) * (len(x_shape) - 1)))
107
+
108
+
109
+ def checkpoint(func, inputs, params, flag):
110
+ """
111
+ Evaluate a function without caching intermediate activations, allowing for
112
+ reduced memory at the expense of extra compute in the backward pass.
113
+ :param func: the function to evaluate.
114
+ :param inputs: the argument sequence to pass to `func`.
115
+ :param params: a sequence of parameters `func` depends on but does not
116
+ explicitly take as arguments.
117
+ :param flag: if False, disable gradient checkpointing.
118
+ """
119
+ if flag:
120
+ args = tuple(inputs) + tuple(params)
121
+ return CheckpointFunction.apply(func, len(inputs), *args)
122
+ else:
123
+ return func(*inputs)
124
+
125
+
126
+ class CheckpointFunction(torch.autograd.Function):
127
+ @staticmethod
128
+ def forward(ctx, run_function, length, *args):
129
+ ctx.run_function = run_function
130
+ ctx.input_tensors = list(args[:length])
131
+ ctx.input_params = list(args[length:])
132
+ ctx.gpu_autocast_kwargs = {"enabled": torch.is_autocast_enabled(),
133
+ "dtype": torch.get_autocast_gpu_dtype(),
134
+ "cache_enabled": torch.is_autocast_cache_enabled()}
135
+ with torch.no_grad():
136
+ output_tensors = ctx.run_function(*ctx.input_tensors)
137
+ return output_tensors
138
+
139
+ @staticmethod
140
+ def backward(ctx, *output_grads):
141
+ ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
142
+ with torch.enable_grad(), \
143
+ torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
144
+ # Fixes a bug where the first op in run_function modifies the
145
+ # Tensor storage in place, which is not allowed for detach()'d
146
+ # Tensors.
147
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
148
+ output_tensors = ctx.run_function(*shallow_copies)
149
+ input_grads = torch.autograd.grad(
150
+ output_tensors,
151
+ ctx.input_tensors + ctx.input_params,
152
+ output_grads,
153
+ allow_unused=True,
154
+ )
155
+ del ctx.input_tensors
156
+ del ctx.input_params
157
+ del output_tensors
158
+ return (None, None) + input_grads
159
+
160
+
161
+ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
162
+ """
163
+ Create sinusoidal timestep embeddings.
164
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
165
+ These may be fractional.
166
+ :param dim: the dimension of the output.
167
+ :param max_period: controls the minimum frequency of the embeddings.
168
+ :return: an [N x dim] Tensor of positional embeddings.
169
+ """
170
+ if not repeat_only:
171
+ half = dim // 2
172
+ freqs = torch.exp(
173
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
174
+ ).to(device=timesteps.device)
175
+ args = timesteps[:, None].float() * freqs[None]
176
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
177
+ if dim % 2:
178
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
179
+ else:
180
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
181
+ return embedding
182
+
183
+
184
+ def zero_module(module):
185
+ """
186
+ Zero out the parameters of a module and return it.
187
+ """
188
+ for p in module.parameters():
189
+ p.detach().zero_()
190
+ return module
191
+
192
+
193
+ def scale_module(module, scale):
194
+ """
195
+ Scale the parameters of a module and return it.
196
+ """
197
+ for p in module.parameters():
198
+ p.detach().mul_(scale)
199
+ return module
200
+
201
+
202
+ def mean_flat(tensor):
203
+ """
204
+ Take the mean over all non-batch dimensions.
205
+ """
206
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
207
+
208
+
209
+ def normalization(channels):
210
+ """
211
+ Make a standard normalization layer.
212
+ :param channels: number of input channels.
213
+ :return: an nn.Module for normalization.
214
+ """
215
+ return GroupNorm32(32, channels)
216
+
217
+
218
+ # PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
219
+ class SiLU(nn.Module):
220
+ def forward(self, x):
221
+ return x * torch.sigmoid(x)
222
+
223
+
224
+ class GroupNorm32(nn.GroupNorm):
225
+ def forward(self, x):
226
+ return super().forward(x.float()).type(x.dtype)
227
+
228
+
229
+ def conv_nd(dims, *args, **kwargs):
230
+ """
231
+ Create a 1D, 2D, or 3D convolution module.
232
+ """
233
+ if dims == 1:
234
+ return nn.Conv1d(*args, **kwargs)
235
+ elif dims == 2:
236
+ return nn.Conv2d(*args, **kwargs)
237
+ elif dims == 3:
238
+ return nn.Conv3d(*args, **kwargs)
239
+ raise ValueError(f"unsupported dimensions: {dims}")
240
+
241
+
242
+ def linear(*args, **kwargs):
243
+ """
244
+ Create a linear module.
245
+ """
246
+ return nn.Linear(*args, **kwargs)
247
+
248
+
249
+ def avg_pool_nd(dims, *args, **kwargs):
250
+ """
251
+ Create a 1D, 2D, or 3D average pooling module.
252
+ """
253
+ if dims == 1:
254
+ return nn.AvgPool1d(*args, **kwargs)
255
+ elif dims == 2:
256
+ return nn.AvgPool2d(*args, **kwargs)
257
+ elif dims == 3:
258
+ return nn.AvgPool3d(*args, **kwargs)
259
+ raise ValueError(f"unsupported dimensions: {dims}")
260
+
261
+
262
+ class HybridConditioner(nn.Module):
263
+
264
+ def __init__(self, c_concat_config, c_crossattn_config):
265
+ super().__init__()
266
+ self.concat_conditioner = instantiate_from_config(c_concat_config)
267
+ self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
268
+
269
+ def forward(self, c_concat, c_crossattn):
270
+ c_concat = self.concat_conditioner(c_concat)
271
+ c_crossattn = self.crossattn_conditioner(c_crossattn)
272
+ return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
273
+
274
+
275
+ def noise_like(shape, device, repeat=False):
276
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
277
+ noise = lambda: torch.randn(shape, device=device)
278
+ return repeat_noise() if repeat else noise()
comfy/ldm/modules/distributions/__init__.py ADDED
File without changes
comfy/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
+ )
comfy/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)
comfy/ldm/modules/encoders/__init__.py ADDED
File without changes
comfy/ldm/modules/encoders/kornia_functions.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+
3
+ from typing import List, Tuple, Union
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+
8
+ #from: https://github.com/kornia/kornia/blob/master/kornia/enhance/normalize.py
9
+
10
+ def enhance_normalize(data: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
11
+ r"""Normalize an image/video tensor with mean and standard deviation.
12
+ .. math::
13
+ \text{input[channel] = (input[channel] - mean[channel]) / std[channel]}
14
+ Where `mean` is :math:`(M_1, ..., M_n)` and `std` :math:`(S_1, ..., S_n)` for `n` channels,
15
+ Args:
16
+ data: Image tensor of size :math:`(B, C, *)`.
17
+ mean: Mean for each channel.
18
+ std: Standard deviations for each channel.
19
+ Return:
20
+ Normalised tensor with same size as input :math:`(B, C, *)`.
21
+ Examples:
22
+ >>> x = torch.rand(1, 4, 3, 3)
23
+ >>> out = normalize(x, torch.tensor([0.0]), torch.tensor([255.]))
24
+ >>> out.shape
25
+ torch.Size([1, 4, 3, 3])
26
+ >>> x = torch.rand(1, 4, 3, 3)
27
+ >>> mean = torch.zeros(4)
28
+ >>> std = 255. * torch.ones(4)
29
+ >>> out = normalize(x, mean, std)
30
+ >>> out.shape
31
+ torch.Size([1, 4, 3, 3])
32
+ """
33
+ shape = data.shape
34
+ if len(mean.shape) == 0 or mean.shape[0] == 1:
35
+ mean = mean.expand(shape[1])
36
+ if len(std.shape) == 0 or std.shape[0] == 1:
37
+ std = std.expand(shape[1])
38
+
39
+ # Allow broadcast on channel dimension
40
+ if mean.shape and mean.shape[0] != 1:
41
+ if mean.shape[0] != data.shape[1] and mean.shape[:2] != data.shape[:2]:
42
+ raise ValueError(f"mean length and number of channels do not match. Got {mean.shape} and {data.shape}.")
43
+
44
+ # Allow broadcast on channel dimension
45
+ if std.shape and std.shape[0] != 1:
46
+ if std.shape[0] != data.shape[1] and std.shape[:2] != data.shape[:2]:
47
+ raise ValueError(f"std length and number of channels do not match. Got {std.shape} and {data.shape}.")
48
+
49
+ mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype)
50
+ std = torch.as_tensor(std, device=data.device, dtype=data.dtype)
51
+
52
+ if mean.shape:
53
+ mean = mean[..., :, None]
54
+ if std.shape:
55
+ std = std[..., :, None]
56
+
57
+ out: torch.Tensor = (data.view(shape[0], shape[1], -1) - mean) / std
58
+
59
+ return out.view(shape)
comfy/ldm/modules/encoders/modules.py ADDED
@@ -0,0 +1,314 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from . import kornia_functions
4
+ from torch.utils.checkpoint import checkpoint
5
+
6
+ from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
7
+
8
+ import open_clip
9
+ from ldm.util import default, count_params
10
+
11
+
12
+ class AbstractEncoder(nn.Module):
13
+ def __init__(self):
14
+ super().__init__()
15
+
16
+ def encode(self, *args, **kwargs):
17
+ raise NotImplementedError
18
+
19
+
20
+ class IdentityEncoder(AbstractEncoder):
21
+
22
+ def encode(self, x):
23
+ return x
24
+
25
+
26
+ class ClassEmbedder(nn.Module):
27
+ def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1):
28
+ super().__init__()
29
+ self.key = key
30
+ self.embedding = nn.Embedding(n_classes, embed_dim)
31
+ self.n_classes = n_classes
32
+ self.ucg_rate = ucg_rate
33
+
34
+ def forward(self, batch, key=None, disable_dropout=False):
35
+ if key is None:
36
+ key = self.key
37
+ # this is for use in crossattn
38
+ c = batch[key][:, None]
39
+ if self.ucg_rate > 0. and not disable_dropout:
40
+ mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
41
+ c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1)
42
+ c = c.long()
43
+ c = self.embedding(c)
44
+ return c
45
+
46
+ def get_unconditional_conditioning(self, bs, device="cuda"):
47
+ uc_class = self.n_classes - 1 # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
48
+ uc = torch.ones((bs,), device=device) * uc_class
49
+ uc = {self.key: uc}
50
+ return uc
51
+
52
+
53
+ def disabled_train(self, mode=True):
54
+ """Overwrite model.train with this function to make sure train/eval mode
55
+ does not change anymore."""
56
+ return self
57
+
58
+
59
+ class FrozenT5Embedder(AbstractEncoder):
60
+ """Uses the T5 transformer encoder for text"""
61
+
62
+ def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77,
63
+ freeze=True): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
64
+ super().__init__()
65
+ self.tokenizer = T5Tokenizer.from_pretrained(version)
66
+ self.transformer = T5EncoderModel.from_pretrained(version)
67
+ self.device = device
68
+ self.max_length = max_length # TODO: typical value?
69
+ if freeze:
70
+ self.freeze()
71
+
72
+ def freeze(self):
73
+ self.transformer = self.transformer.eval()
74
+ # self.train = disabled_train
75
+ for param in self.parameters():
76
+ param.requires_grad = False
77
+
78
+ def forward(self, text):
79
+ batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
80
+ return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
81
+ tokens = batch_encoding["input_ids"].to(self.device)
82
+ outputs = self.transformer(input_ids=tokens)
83
+
84
+ z = outputs.last_hidden_state
85
+ return z
86
+
87
+ def encode(self, text):
88
+ return self(text)
89
+
90
+
91
+ class FrozenCLIPEmbedder(AbstractEncoder):
92
+ """Uses the CLIP transformer encoder for text (from huggingface)"""
93
+ LAYERS = [
94
+ "last",
95
+ "pooled",
96
+ "hidden"
97
+ ]
98
+
99
+ def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
100
+ freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32
101
+ super().__init__()
102
+ assert layer in self.LAYERS
103
+ self.tokenizer = CLIPTokenizer.from_pretrained(version)
104
+ self.transformer = CLIPTextModel.from_pretrained(version)
105
+ self.device = device
106
+ self.max_length = max_length
107
+ if freeze:
108
+ self.freeze()
109
+ self.layer = layer
110
+ self.layer_idx = layer_idx
111
+ if layer == "hidden":
112
+ assert layer_idx is not None
113
+ assert 0 <= abs(layer_idx) <= 12
114
+
115
+ def freeze(self):
116
+ self.transformer = self.transformer.eval()
117
+ # self.train = disabled_train
118
+ for param in self.parameters():
119
+ param.requires_grad = False
120
+
121
+ def forward(self, text):
122
+ batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
123
+ return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
124
+ tokens = batch_encoding["input_ids"].to(self.device)
125
+ outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden")
126
+ if self.layer == "last":
127
+ z = outputs.last_hidden_state
128
+ elif self.layer == "pooled":
129
+ z = outputs.pooler_output[:, None, :]
130
+ else:
131
+ z = outputs.hidden_states[self.layer_idx]
132
+ return z
133
+
134
+ def encode(self, text):
135
+ return self(text)
136
+
137
+
138
+ class ClipImageEmbedder(nn.Module):
139
+ def __init__(
140
+ self,
141
+ model,
142
+ jit=False,
143
+ device='cuda' if torch.cuda.is_available() else 'cpu',
144
+ antialias=True,
145
+ ucg_rate=0.
146
+ ):
147
+ super().__init__()
148
+ from clip import load as load_clip
149
+ self.model, _ = load_clip(name=model, device=device, jit=jit)
150
+
151
+ self.antialias = antialias
152
+
153
+ self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
154
+ self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
155
+ self.ucg_rate = ucg_rate
156
+
157
+ def preprocess(self, x):
158
+ # normalize to [0,1]
159
+ # x = kornia_functions.geometry_resize(x, (224, 224),
160
+ # interpolation='bicubic', align_corners=True,
161
+ # antialias=self.antialias)
162
+ x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bicubic', align_corners=True, antialias=True)
163
+ x = (x + 1.) / 2.
164
+ # re-normalize according to clip
165
+ x = kornia_functions.enhance_normalize(x, self.mean, self.std)
166
+ return x
167
+
168
+ def forward(self, x, no_dropout=False):
169
+ # x is assumed to be in range [-1,1]
170
+ out = self.model.encode_image(self.preprocess(x))
171
+ out = out.to(x.dtype)
172
+ if self.ucg_rate > 0. and not no_dropout:
173
+ out = torch.bernoulli((1. - self.ucg_rate) * torch.ones(out.shape[0], device=out.device))[:, None] * out
174
+ return out
175
+
176
+
177
+ class FrozenOpenCLIPEmbedder(AbstractEncoder):
178
+ """
179
+ Uses the OpenCLIP transformer encoder for text
180
+ """
181
+ LAYERS = [
182
+ # "pooled",
183
+ "last",
184
+ "penultimate"
185
+ ]
186
+
187
+ def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
188
+ freeze=True, layer="last"):
189
+ super().__init__()
190
+ assert layer in self.LAYERS
191
+ model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)
192
+ del model.visual
193
+ self.model = model
194
+
195
+ self.device = device
196
+ self.max_length = max_length
197
+ if freeze:
198
+ self.freeze()
199
+ self.layer = layer
200
+ if self.layer == "last":
201
+ self.layer_idx = 0
202
+ elif self.layer == "penultimate":
203
+ self.layer_idx = 1
204
+ else:
205
+ raise NotImplementedError()
206
+
207
+ def freeze(self):
208
+ self.model = self.model.eval()
209
+ for param in self.parameters():
210
+ param.requires_grad = False
211
+
212
+ def forward(self, text):
213
+ tokens = open_clip.tokenize(text)
214
+ z = self.encode_with_transformer(tokens.to(self.device))
215
+ return z
216
+
217
+ def encode_with_transformer(self, text):
218
+ x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model]
219
+ x = x + self.model.positional_embedding
220
+ x = x.permute(1, 0, 2) # NLD -> LND
221
+ x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
222
+ x = x.permute(1, 0, 2) # LND -> NLD
223
+ x = self.model.ln_final(x)
224
+ return x
225
+
226
+ def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
227
+ for i, r in enumerate(self.model.transformer.resblocks):
228
+ if i == len(self.model.transformer.resblocks) - self.layer_idx:
229
+ break
230
+ if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
231
+ x = checkpoint(r, x, attn_mask)
232
+ else:
233
+ x = r(x, attn_mask=attn_mask)
234
+ return x
235
+
236
+ def encode(self, text):
237
+ return self(text)
238
+
239
+
240
+ class FrozenOpenCLIPImageEmbedder(AbstractEncoder):
241
+ """
242
+ Uses the OpenCLIP vision transformer encoder for images
243
+ """
244
+
245
+ def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
246
+ freeze=True, layer="pooled", antialias=True, ucg_rate=0.):
247
+ super().__init__()
248
+ model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'),
249
+ pretrained=version, )
250
+ del model.transformer
251
+ self.model = model
252
+
253
+ self.device = device
254
+ self.max_length = max_length
255
+ if freeze:
256
+ self.freeze()
257
+ self.layer = layer
258
+ if self.layer == "penultimate":
259
+ raise NotImplementedError()
260
+ self.layer_idx = 1
261
+
262
+ self.antialias = antialias
263
+
264
+ self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
265
+ self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
266
+ self.ucg_rate = ucg_rate
267
+
268
+ def preprocess(self, x):
269
+ # normalize to [0,1]
270
+ # x = kornia.geometry.resize(x, (224, 224),
271
+ # interpolation='bicubic', align_corners=True,
272
+ # antialias=self.antialias)
273
+ x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bicubic', align_corners=True, antialias=True)
274
+ x = (x + 1.) / 2.
275
+ # renormalize according to clip
276
+ x = kornia_functions.enhance_normalize(x, self.mean, self.std)
277
+ return x
278
+
279
+ def freeze(self):
280
+ self.model = self.model.eval()
281
+ for param in self.parameters():
282
+ param.requires_grad = False
283
+
284
+ def forward(self, image, no_dropout=False):
285
+ z = self.encode_with_vision_transformer(image)
286
+ if self.ucg_rate > 0. and not no_dropout:
287
+ z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z
288
+ return z
289
+
290
+ def encode_with_vision_transformer(self, img):
291
+ img = self.preprocess(img)
292
+ x = self.model.visual(img)
293
+ return x
294
+
295
+ def encode(self, text):
296
+ return self(text)
297
+
298
+
299
+ class FrozenCLIPT5Encoder(AbstractEncoder):
300
+ def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda",
301
+ clip_max_length=77, t5_max_length=77):
302
+ super().__init__()
303
+ self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)
304
+ self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length)
305
+ print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder) * 1.e-6:.2f} M parameters, "
306
+ f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder) * 1.e-6:.2f} M params.")
307
+
308
+ def encode(self, text):
309
+ return self(text)
310
+
311
+ def forward(self, text):
312
+ clip_z = self.clip_encoder.encode(text)
313
+ t5_z = self.t5_encoder.encode(text)
314
+ return [clip_z, t5_z]
comfy/ldm/modules/encoders/noise_aug_modules.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
2
+ from ldm.modules.diffusionmodules.openaimodel import Timestep
3
+ import torch
4
+
5
+ class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
6
+ def __init__(self, *args, clip_stats_path=None, timestep_dim=256, **kwargs):
7
+ super().__init__(*args, **kwargs)
8
+ if clip_stats_path is None:
9
+ clip_mean, clip_std = torch.zeros(timestep_dim), torch.ones(timestep_dim)
10
+ else:
11
+ clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu")
12
+ self.register_buffer("data_mean", clip_mean[None, :], persistent=False)
13
+ self.register_buffer("data_std", clip_std[None, :], persistent=False)
14
+ self.time_embed = Timestep(timestep_dim)
15
+
16
+ def scale(self, x):
17
+ # re-normalize to centered mean and unit variance
18
+ x = (x - self.data_mean) * 1. / self.data_std
19
+ return x
20
+
21
+ def unscale(self, x):
22
+ # back to original data stats
23
+ x = (x * self.data_std) + self.data_mean
24
+ return x
25
+
26
+ def forward(self, x, noise_level=None):
27
+ if noise_level is None:
28
+ noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
29
+ else:
30
+ assert isinstance(noise_level, torch.Tensor)
31
+ x = self.scale(x)
32
+ z = self.q_sample(x, noise_level)
33
+ z = self.unscale(z)
34
+ noise_level = self.time_embed(noise_level)
35
+ return z, noise_level
comfy/ldm/modules/image_degradation/__init__.py ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ from ldm.modules.image_degradation.bsrgan import degradation_bsrgan_variant as degradation_fn_bsr
2
+ from ldm.modules.image_degradation.bsrgan_light import degradation_bsrgan_variant as degradation_fn_bsr_light
comfy/ldm/modules/image_degradation/bsrgan.py ADDED
@@ -0,0 +1,730 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+ """
3
+ # --------------------------------------------
4
+ # Super-Resolution
5
+ # --------------------------------------------
6
+ #
7
+ # Kai Zhang (cskaizhang@gmail.com)
8
+ # https://github.com/cszn
9
+ # From 2019/03--2021/08
10
+ # --------------------------------------------
11
+ """
12
+
13
+ import numpy as np
14
+ import cv2
15
+ import torch
16
+
17
+ from functools import partial
18
+ import random
19
+ from scipy import ndimage
20
+ import scipy
21
+ import scipy.stats as ss
22
+ from scipy.interpolate import interp2d
23
+ from scipy.linalg import orth
24
+ import albumentations
25
+
26
+ import ldm.modules.image_degradation.utils_image as util
27
+
28
+
29
+ def modcrop_np(img, sf):
30
+ '''
31
+ Args:
32
+ img: numpy image, WxH or WxHxC
33
+ sf: scale factor
34
+ Return:
35
+ cropped image
36
+ '''
37
+ w, h = img.shape[:2]
38
+ im = np.copy(img)
39
+ return im[:w - w % sf, :h - h % sf, ...]
40
+
41
+
42
+ """
43
+ # --------------------------------------------
44
+ # anisotropic Gaussian kernels
45
+ # --------------------------------------------
46
+ """
47
+
48
+
49
+ def analytic_kernel(k):
50
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
51
+ k_size = k.shape[0]
52
+ # Calculate the big kernels size
53
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
54
+ # Loop over the small kernel to fill the big one
55
+ for r in range(k_size):
56
+ for c in range(k_size):
57
+ big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
58
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
59
+ crop = k_size // 2
60
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
61
+ # Normalize to 1
62
+ return cropped_big_k / cropped_big_k.sum()
63
+
64
+
65
+ def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
66
+ """ generate an anisotropic Gaussian kernel
67
+ Args:
68
+ ksize : e.g., 15, kernel size
69
+ theta : [0, pi], rotation angle range
70
+ l1 : [0.1,50], scaling of eigenvalues
71
+ l2 : [0.1,l1], scaling of eigenvalues
72
+ If l1 = l2, will get an isotropic Gaussian kernel.
73
+ Returns:
74
+ k : kernel
75
+ """
76
+
77
+ v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
78
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
79
+ D = np.array([[l1, 0], [0, l2]])
80
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
81
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
82
+
83
+ return k
84
+
85
+
86
+ def gm_blur_kernel(mean, cov, size=15):
87
+ center = size / 2.0 + 0.5
88
+ k = np.zeros([size, size])
89
+ for y in range(size):
90
+ for x in range(size):
91
+ cy = y - center + 1
92
+ cx = x - center + 1
93
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
94
+
95
+ k = k / np.sum(k)
96
+ return k
97
+
98
+
99
+ def shift_pixel(x, sf, upper_left=True):
100
+ """shift pixel for super-resolution with different scale factors
101
+ Args:
102
+ x: WxHxC or WxH
103
+ sf: scale factor
104
+ upper_left: shift direction
105
+ """
106
+ h, w = x.shape[:2]
107
+ shift = (sf - 1) * 0.5
108
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
109
+ if upper_left:
110
+ x1 = xv + shift
111
+ y1 = yv + shift
112
+ else:
113
+ x1 = xv - shift
114
+ y1 = yv - shift
115
+
116
+ x1 = np.clip(x1, 0, w - 1)
117
+ y1 = np.clip(y1, 0, h - 1)
118
+
119
+ if x.ndim == 2:
120
+ x = interp2d(xv, yv, x)(x1, y1)
121
+ if x.ndim == 3:
122
+ for i in range(x.shape[-1]):
123
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
124
+
125
+ return x
126
+
127
+
128
+ def blur(x, k):
129
+ '''
130
+ x: image, NxcxHxW
131
+ k: kernel, Nx1xhxw
132
+ '''
133
+ n, c = x.shape[:2]
134
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
135
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
136
+ k = k.repeat(1, c, 1, 1)
137
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
138
+ x = x.view(1, -1, x.shape[2], x.shape[3])
139
+ x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c)
140
+ x = x.view(n, c, x.shape[2], x.shape[3])
141
+
142
+ return x
143
+
144
+
145
+ def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
146
+ """"
147
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
148
+ # Kai Zhang
149
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
150
+ # max_var = 2.5 * sf
151
+ """
152
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
153
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
154
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
155
+ theta = np.random.rand() * np.pi # random theta
156
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
157
+
158
+ # Set COV matrix using Lambdas and Theta
159
+ LAMBDA = np.diag([lambda_1, lambda_2])
160
+ Q = np.array([[np.cos(theta), -np.sin(theta)],
161
+ [np.sin(theta), np.cos(theta)]])
162
+ SIGMA = Q @ LAMBDA @ Q.T
163
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
164
+
165
+ # Set expectation position (shifting kernel for aligned image)
166
+ MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
167
+ MU = MU[None, None, :, None]
168
+
169
+ # Create meshgrid for Gaussian
170
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
171
+ Z = np.stack([X, Y], 2)[:, :, :, None]
172
+
173
+ # Calcualte Gaussian for every pixel of the kernel
174
+ ZZ = Z - MU
175
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
176
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
177
+
178
+ # shift the kernel so it will be centered
179
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
180
+
181
+ # Normalize the kernel and return
182
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
183
+ kernel = raw_kernel / np.sum(raw_kernel)
184
+ return kernel
185
+
186
+
187
+ def fspecial_gaussian(hsize, sigma):
188
+ hsize = [hsize, hsize]
189
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
190
+ std = sigma
191
+ [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
192
+ arg = -(x * x + y * y) / (2 * std * std)
193
+ h = np.exp(arg)
194
+ h[h < scipy.finfo(float).eps * h.max()] = 0
195
+ sumh = h.sum()
196
+ if sumh != 0:
197
+ h = h / sumh
198
+ return h
199
+
200
+
201
+ def fspecial_laplacian(alpha):
202
+ alpha = max([0, min([alpha, 1])])
203
+ h1 = alpha / (alpha + 1)
204
+ h2 = (1 - alpha) / (alpha + 1)
205
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
206
+ h = np.array(h)
207
+ return h
208
+
209
+
210
+ def fspecial(filter_type, *args, **kwargs):
211
+ '''
212
+ python code from:
213
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
214
+ '''
215
+ if filter_type == 'gaussian':
216
+ return fspecial_gaussian(*args, **kwargs)
217
+ if filter_type == 'laplacian':
218
+ return fspecial_laplacian(*args, **kwargs)
219
+
220
+
221
+ """
222
+ # --------------------------------------------
223
+ # degradation models
224
+ # --------------------------------------------
225
+ """
226
+
227
+
228
+ def bicubic_degradation(x, sf=3):
229
+ '''
230
+ Args:
231
+ x: HxWxC image, [0, 1]
232
+ sf: down-scale factor
233
+ Return:
234
+ bicubicly downsampled LR image
235
+ '''
236
+ x = util.imresize_np(x, scale=1 / sf)
237
+ return x
238
+
239
+
240
+ def srmd_degradation(x, k, sf=3):
241
+ ''' blur + bicubic downsampling
242
+ Args:
243
+ x: HxWxC image, [0, 1]
244
+ k: hxw, double
245
+ sf: down-scale factor
246
+ Return:
247
+ downsampled LR image
248
+ Reference:
249
+ @inproceedings{zhang2018learning,
250
+ title={Learning a single convolutional super-resolution network for multiple degradations},
251
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
252
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
253
+ pages={3262--3271},
254
+ year={2018}
255
+ }
256
+ '''
257
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror'
258
+ x = bicubic_degradation(x, sf=sf)
259
+ return x
260
+
261
+
262
+ def dpsr_degradation(x, k, sf=3):
263
+ ''' bicubic downsampling + blur
264
+ Args:
265
+ x: HxWxC image, [0, 1]
266
+ k: hxw, double
267
+ sf: down-scale factor
268
+ Return:
269
+ downsampled LR image
270
+ Reference:
271
+ @inproceedings{zhang2019deep,
272
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
273
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
274
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
275
+ pages={1671--1681},
276
+ year={2019}
277
+ }
278
+ '''
279
+ x = bicubic_degradation(x, sf=sf)
280
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
281
+ return x
282
+
283
+
284
+ def classical_degradation(x, k, sf=3):
285
+ ''' blur + downsampling
286
+ Args:
287
+ x: HxWxC image, [0, 1]/[0, 255]
288
+ k: hxw, double
289
+ sf: down-scale factor
290
+ Return:
291
+ downsampled LR image
292
+ '''
293
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
294
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
295
+ st = 0
296
+ return x[st::sf, st::sf, ...]
297
+
298
+
299
+ def add_sharpening(img, weight=0.5, radius=50, threshold=10):
300
+ """USM sharpening. borrowed from real-ESRGAN
301
+ Input image: I; Blurry image: B.
302
+ 1. K = I + weight * (I - B)
303
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
304
+ 3. Blur mask:
305
+ 4. Out = Mask * K + (1 - Mask) * I
306
+ Args:
307
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
308
+ weight (float): Sharp weight. Default: 1.
309
+ radius (float): Kernel size of Gaussian blur. Default: 50.
310
+ threshold (int):
311
+ """
312
+ if radius % 2 == 0:
313
+ radius += 1
314
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
315
+ residual = img - blur
316
+ mask = np.abs(residual) * 255 > threshold
317
+ mask = mask.astype('float32')
318
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
319
+
320
+ K = img + weight * residual
321
+ K = np.clip(K, 0, 1)
322
+ return soft_mask * K + (1 - soft_mask) * img
323
+
324
+
325
+ def add_blur(img, sf=4):
326
+ wd2 = 4.0 + sf
327
+ wd = 2.0 + 0.2 * sf
328
+ if random.random() < 0.5:
329
+ l1 = wd2 * random.random()
330
+ l2 = wd2 * random.random()
331
+ k = anisotropic_Gaussian(ksize=2 * random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2)
332
+ else:
333
+ k = fspecial('gaussian', 2 * random.randint(2, 11) + 3, wd * random.random())
334
+ img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
335
+
336
+ return img
337
+
338
+
339
+ def add_resize(img, sf=4):
340
+ rnum = np.random.rand()
341
+ if rnum > 0.8: # up
342
+ sf1 = random.uniform(1, 2)
343
+ elif rnum < 0.7: # down
344
+ sf1 = random.uniform(0.5 / sf, 1)
345
+ else:
346
+ sf1 = 1.0
347
+ img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3]))
348
+ img = np.clip(img, 0.0, 1.0)
349
+
350
+ return img
351
+
352
+
353
+ # def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
354
+ # noise_level = random.randint(noise_level1, noise_level2)
355
+ # rnum = np.random.rand()
356
+ # if rnum > 0.6: # add color Gaussian noise
357
+ # img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
358
+ # elif rnum < 0.4: # add grayscale Gaussian noise
359
+ # img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
360
+ # else: # add noise
361
+ # L = noise_level2 / 255.
362
+ # D = np.diag(np.random.rand(3))
363
+ # U = orth(np.random.rand(3, 3))
364
+ # conv = np.dot(np.dot(np.transpose(U), D), U)
365
+ # img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
366
+ # img = np.clip(img, 0.0, 1.0)
367
+ # return img
368
+
369
+ def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
370
+ noise_level = random.randint(noise_level1, noise_level2)
371
+ rnum = np.random.rand()
372
+ if rnum > 0.6: # add color Gaussian noise
373
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
374
+ elif rnum < 0.4: # add grayscale Gaussian noise
375
+ img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
376
+ else: # add noise
377
+ L = noise_level2 / 255.
378
+ D = np.diag(np.random.rand(3))
379
+ U = orth(np.random.rand(3, 3))
380
+ conv = np.dot(np.dot(np.transpose(U), D), U)
381
+ img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
382
+ img = np.clip(img, 0.0, 1.0)
383
+ return img
384
+
385
+
386
+ def add_speckle_noise(img, noise_level1=2, noise_level2=25):
387
+ noise_level = random.randint(noise_level1, noise_level2)
388
+ img = np.clip(img, 0.0, 1.0)
389
+ rnum = random.random()
390
+ if rnum > 0.6:
391
+ img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
392
+ elif rnum < 0.4:
393
+ img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
394
+ else:
395
+ L = noise_level2 / 255.
396
+ D = np.diag(np.random.rand(3))
397
+ U = orth(np.random.rand(3, 3))
398
+ conv = np.dot(np.dot(np.transpose(U), D), U)
399
+ img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
400
+ img = np.clip(img, 0.0, 1.0)
401
+ return img
402
+
403
+
404
+ def add_Poisson_noise(img):
405
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.
406
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
407
+ if random.random() < 0.5:
408
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
409
+ else:
410
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
411
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
412
+ noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
413
+ img += noise_gray[:, :, np.newaxis]
414
+ img = np.clip(img, 0.0, 1.0)
415
+ return img
416
+
417
+
418
+ def add_JPEG_noise(img):
419
+ quality_factor = random.randint(30, 95)
420
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
421
+ result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
422
+ img = cv2.imdecode(encimg, 1)
423
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
424
+ return img
425
+
426
+
427
+ def random_crop(lq, hq, sf=4, lq_patchsize=64):
428
+ h, w = lq.shape[:2]
429
+ rnd_h = random.randint(0, h - lq_patchsize)
430
+ rnd_w = random.randint(0, w - lq_patchsize)
431
+ lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
432
+
433
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
434
+ hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
435
+ return lq, hq
436
+
437
+
438
+ def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
439
+ """
440
+ This is the degradation model of BSRGAN from the paper
441
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
442
+ ----------
443
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
444
+ sf: scale factor
445
+ isp_model: camera ISP model
446
+ Returns
447
+ -------
448
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
449
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
450
+ """
451
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
452
+ sf_ori = sf
453
+
454
+ h1, w1 = img.shape[:2]
455
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
456
+ h, w = img.shape[:2]
457
+
458
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
459
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
460
+
461
+ hq = img.copy()
462
+
463
+ if sf == 4 and random.random() < scale2_prob: # downsample1
464
+ if np.random.rand() < 0.5:
465
+ img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
466
+ interpolation=random.choice([1, 2, 3]))
467
+ else:
468
+ img = util.imresize_np(img, 1 / 2, True)
469
+ img = np.clip(img, 0.0, 1.0)
470
+ sf = 2
471
+
472
+ shuffle_order = random.sample(range(7), 7)
473
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
474
+ if idx1 > idx2: # keep downsample3 last
475
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
476
+
477
+ for i in shuffle_order:
478
+
479
+ if i == 0:
480
+ img = add_blur(img, sf=sf)
481
+
482
+ elif i == 1:
483
+ img = add_blur(img, sf=sf)
484
+
485
+ elif i == 2:
486
+ a, b = img.shape[1], img.shape[0]
487
+ # downsample2
488
+ if random.random() < 0.75:
489
+ sf1 = random.uniform(1, 2 * sf)
490
+ img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
491
+ interpolation=random.choice([1, 2, 3]))
492
+ else:
493
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
494
+ k_shifted = shift_pixel(k, sf)
495
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
496
+ img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
497
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
498
+ img = np.clip(img, 0.0, 1.0)
499
+
500
+ elif i == 3:
501
+ # downsample3
502
+ img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
503
+ img = np.clip(img, 0.0, 1.0)
504
+
505
+ elif i == 4:
506
+ # add Gaussian noise
507
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
508
+
509
+ elif i == 5:
510
+ # add JPEG noise
511
+ if random.random() < jpeg_prob:
512
+ img = add_JPEG_noise(img)
513
+
514
+ elif i == 6:
515
+ # add processed camera sensor noise
516
+ if random.random() < isp_prob and isp_model is not None:
517
+ with torch.no_grad():
518
+ img, hq = isp_model.forward(img.copy(), hq)
519
+
520
+ # add final JPEG compression noise
521
+ img = add_JPEG_noise(img)
522
+
523
+ # random crop
524
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
525
+
526
+ return img, hq
527
+
528
+
529
+ # todo no isp_model?
530
+ def degradation_bsrgan_variant(image, sf=4, isp_model=None):
531
+ """
532
+ This is the degradation model of BSRGAN from the paper
533
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
534
+ ----------
535
+ sf: scale factor
536
+ isp_model: camera ISP model
537
+ Returns
538
+ -------
539
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
540
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
541
+ """
542
+ image = util.uint2single(image)
543
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
544
+ sf_ori = sf
545
+
546
+ h1, w1 = image.shape[:2]
547
+ image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
548
+ h, w = image.shape[:2]
549
+
550
+ hq = image.copy()
551
+
552
+ if sf == 4 and random.random() < scale2_prob: # downsample1
553
+ if np.random.rand() < 0.5:
554
+ image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
555
+ interpolation=random.choice([1, 2, 3]))
556
+ else:
557
+ image = util.imresize_np(image, 1 / 2, True)
558
+ image = np.clip(image, 0.0, 1.0)
559
+ sf = 2
560
+
561
+ shuffle_order = random.sample(range(7), 7)
562
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
563
+ if idx1 > idx2: # keep downsample3 last
564
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
565
+
566
+ for i in shuffle_order:
567
+
568
+ if i == 0:
569
+ image = add_blur(image, sf=sf)
570
+
571
+ elif i == 1:
572
+ image = add_blur(image, sf=sf)
573
+
574
+ elif i == 2:
575
+ a, b = image.shape[1], image.shape[0]
576
+ # downsample2
577
+ if random.random() < 0.75:
578
+ sf1 = random.uniform(1, 2 * sf)
579
+ image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])),
580
+ interpolation=random.choice([1, 2, 3]))
581
+ else:
582
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
583
+ k_shifted = shift_pixel(k, sf)
584
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
585
+ image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror')
586
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
587
+ image = np.clip(image, 0.0, 1.0)
588
+
589
+ elif i == 3:
590
+ # downsample3
591
+ image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
592
+ image = np.clip(image, 0.0, 1.0)
593
+
594
+ elif i == 4:
595
+ # add Gaussian noise
596
+ image = add_Gaussian_noise(image, noise_level1=2, noise_level2=25)
597
+
598
+ elif i == 5:
599
+ # add JPEG noise
600
+ if random.random() < jpeg_prob:
601
+ image = add_JPEG_noise(image)
602
+
603
+ # elif i == 6:
604
+ # # add processed camera sensor noise
605
+ # if random.random() < isp_prob and isp_model is not None:
606
+ # with torch.no_grad():
607
+ # img, hq = isp_model.forward(img.copy(), hq)
608
+
609
+ # add final JPEG compression noise
610
+ image = add_JPEG_noise(image)
611
+ image = util.single2uint(image)
612
+ example = {"image":image}
613
+ return example
614
+
615
+
616
+ # TODO incase there is a pickle error one needs to replace a += x with a = a + x in add_speckle_noise etc...
617
+ def degradation_bsrgan_plus(img, sf=4, shuffle_prob=0.5, use_sharp=True, lq_patchsize=64, isp_model=None):
618
+ """
619
+ This is an extended degradation model by combining
620
+ the degradation models of BSRGAN and Real-ESRGAN
621
+ ----------
622
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
623
+ sf: scale factor
624
+ use_shuffle: the degradation shuffle
625
+ use_sharp: sharpening the img
626
+ Returns
627
+ -------
628
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
629
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
630
+ """
631
+
632
+ h1, w1 = img.shape[:2]
633
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
634
+ h, w = img.shape[:2]
635
+
636
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
637
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
638
+
639
+ if use_sharp:
640
+ img = add_sharpening(img)
641
+ hq = img.copy()
642
+
643
+ if random.random() < shuffle_prob:
644
+ shuffle_order = random.sample(range(13), 13)
645
+ else:
646
+ shuffle_order = list(range(13))
647
+ # local shuffle for noise, JPEG is always the last one
648
+ shuffle_order[2:6] = random.sample(shuffle_order[2:6], len(range(2, 6)))
649
+ shuffle_order[9:13] = random.sample(shuffle_order[9:13], len(range(9, 13)))
650
+
651
+ poisson_prob, speckle_prob, isp_prob = 0.1, 0.1, 0.1
652
+
653
+ for i in shuffle_order:
654
+ if i == 0:
655
+ img = add_blur(img, sf=sf)
656
+ elif i == 1:
657
+ img = add_resize(img, sf=sf)
658
+ elif i == 2:
659
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
660
+ elif i == 3:
661
+ if random.random() < poisson_prob:
662
+ img = add_Poisson_noise(img)
663
+ elif i == 4:
664
+ if random.random() < speckle_prob:
665
+ img = add_speckle_noise(img)
666
+ elif i == 5:
667
+ if random.random() < isp_prob and isp_model is not None:
668
+ with torch.no_grad():
669
+ img, hq = isp_model.forward(img.copy(), hq)
670
+ elif i == 6:
671
+ img = add_JPEG_noise(img)
672
+ elif i == 7:
673
+ img = add_blur(img, sf=sf)
674
+ elif i == 8:
675
+ img = add_resize(img, sf=sf)
676
+ elif i == 9:
677
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
678
+ elif i == 10:
679
+ if random.random() < poisson_prob:
680
+ img = add_Poisson_noise(img)
681
+ elif i == 11:
682
+ if random.random() < speckle_prob:
683
+ img = add_speckle_noise(img)
684
+ elif i == 12:
685
+ if random.random() < isp_prob and isp_model is not None:
686
+ with torch.no_grad():
687
+ img, hq = isp_model.forward(img.copy(), hq)
688
+ else:
689
+ print('check the shuffle!')
690
+
691
+ # resize to desired size
692
+ img = cv2.resize(img, (int(1 / sf * hq.shape[1]), int(1 / sf * hq.shape[0])),
693
+ interpolation=random.choice([1, 2, 3]))
694
+
695
+ # add final JPEG compression noise
696
+ img = add_JPEG_noise(img)
697
+
698
+ # random crop
699
+ img, hq = random_crop(img, hq, sf, lq_patchsize)
700
+
701
+ return img, hq
702
+
703
+
704
+ if __name__ == '__main__':
705
+ print("hey")
706
+ img = util.imread_uint('utils/test.png', 3)
707
+ print(img)
708
+ img = util.uint2single(img)
709
+ print(img)
710
+ img = img[:448, :448]
711
+ h = img.shape[0] // 4
712
+ print("resizing to", h)
713
+ sf = 4
714
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
715
+ for i in range(20):
716
+ print(i)
717
+ img_lq = deg_fn(img)
718
+ print(img_lq)
719
+ img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img)["image"]
720
+ print(img_lq.shape)
721
+ print("bicubic", img_lq_bicubic.shape)
722
+ print(img_hq.shape)
723
+ lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
724
+ interpolation=0)
725
+ lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
726
+ interpolation=0)
727
+ img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1)
728
+ util.imsave(img_concat, str(i) + '.png')
729
+
730
+
comfy/ldm/modules/image_degradation/bsrgan_light.py ADDED
@@ -0,0 +1,651 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+ import numpy as np
3
+ import cv2
4
+ import torch
5
+
6
+ from functools import partial
7
+ import random
8
+ from scipy import ndimage
9
+ import scipy
10
+ import scipy.stats as ss
11
+ from scipy.interpolate import interp2d
12
+ from scipy.linalg import orth
13
+ import albumentations
14
+
15
+ import ldm.modules.image_degradation.utils_image as util
16
+
17
+ """
18
+ # --------------------------------------------
19
+ # Super-Resolution
20
+ # --------------------------------------------
21
+ #
22
+ # Kai Zhang (cskaizhang@gmail.com)
23
+ # https://github.com/cszn
24
+ # From 2019/03--2021/08
25
+ # --------------------------------------------
26
+ """
27
+
28
+ def modcrop_np(img, sf):
29
+ '''
30
+ Args:
31
+ img: numpy image, WxH or WxHxC
32
+ sf: scale factor
33
+ Return:
34
+ cropped image
35
+ '''
36
+ w, h = img.shape[:2]
37
+ im = np.copy(img)
38
+ return im[:w - w % sf, :h - h % sf, ...]
39
+
40
+
41
+ """
42
+ # --------------------------------------------
43
+ # anisotropic Gaussian kernels
44
+ # --------------------------------------------
45
+ """
46
+
47
+
48
+ def analytic_kernel(k):
49
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
50
+ k_size = k.shape[0]
51
+ # Calculate the big kernels size
52
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
53
+ # Loop over the small kernel to fill the big one
54
+ for r in range(k_size):
55
+ for c in range(k_size):
56
+ big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
57
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
58
+ crop = k_size // 2
59
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
60
+ # Normalize to 1
61
+ return cropped_big_k / cropped_big_k.sum()
62
+
63
+
64
+ def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
65
+ """ generate an anisotropic Gaussian kernel
66
+ Args:
67
+ ksize : e.g., 15, kernel size
68
+ theta : [0, pi], rotation angle range
69
+ l1 : [0.1,50], scaling of eigenvalues
70
+ l2 : [0.1,l1], scaling of eigenvalues
71
+ If l1 = l2, will get an isotropic Gaussian kernel.
72
+ Returns:
73
+ k : kernel
74
+ """
75
+
76
+ v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
77
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
78
+ D = np.array([[l1, 0], [0, l2]])
79
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
80
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
81
+
82
+ return k
83
+
84
+
85
+ def gm_blur_kernel(mean, cov, size=15):
86
+ center = size / 2.0 + 0.5
87
+ k = np.zeros([size, size])
88
+ for y in range(size):
89
+ for x in range(size):
90
+ cy = y - center + 1
91
+ cx = x - center + 1
92
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
93
+
94
+ k = k / np.sum(k)
95
+ return k
96
+
97
+
98
+ def shift_pixel(x, sf, upper_left=True):
99
+ """shift pixel for super-resolution with different scale factors
100
+ Args:
101
+ x: WxHxC or WxH
102
+ sf: scale factor
103
+ upper_left: shift direction
104
+ """
105
+ h, w = x.shape[:2]
106
+ shift = (sf - 1) * 0.5
107
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
108
+ if upper_left:
109
+ x1 = xv + shift
110
+ y1 = yv + shift
111
+ else:
112
+ x1 = xv - shift
113
+ y1 = yv - shift
114
+
115
+ x1 = np.clip(x1, 0, w - 1)
116
+ y1 = np.clip(y1, 0, h - 1)
117
+
118
+ if x.ndim == 2:
119
+ x = interp2d(xv, yv, x)(x1, y1)
120
+ if x.ndim == 3:
121
+ for i in range(x.shape[-1]):
122
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
123
+
124
+ return x
125
+
126
+
127
+ def blur(x, k):
128
+ '''
129
+ x: image, NxcxHxW
130
+ k: kernel, Nx1xhxw
131
+ '''
132
+ n, c = x.shape[:2]
133
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
134
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
135
+ k = k.repeat(1, c, 1, 1)
136
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
137
+ x = x.view(1, -1, x.shape[2], x.shape[3])
138
+ x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c)
139
+ x = x.view(n, c, x.shape[2], x.shape[3])
140
+
141
+ return x
142
+
143
+
144
+ def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
145
+ """"
146
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
147
+ # Kai Zhang
148
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
149
+ # max_var = 2.5 * sf
150
+ """
151
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
152
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
153
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
154
+ theta = np.random.rand() * np.pi # random theta
155
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
156
+
157
+ # Set COV matrix using Lambdas and Theta
158
+ LAMBDA = np.diag([lambda_1, lambda_2])
159
+ Q = np.array([[np.cos(theta), -np.sin(theta)],
160
+ [np.sin(theta), np.cos(theta)]])
161
+ SIGMA = Q @ LAMBDA @ Q.T
162
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
163
+
164
+ # Set expectation position (shifting kernel for aligned image)
165
+ MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
166
+ MU = MU[None, None, :, None]
167
+
168
+ # Create meshgrid for Gaussian
169
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
170
+ Z = np.stack([X, Y], 2)[:, :, :, None]
171
+
172
+ # Calcualte Gaussian for every pixel of the kernel
173
+ ZZ = Z - MU
174
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
175
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
176
+
177
+ # shift the kernel so it will be centered
178
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
179
+
180
+ # Normalize the kernel and return
181
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
182
+ kernel = raw_kernel / np.sum(raw_kernel)
183
+ return kernel
184
+
185
+
186
+ def fspecial_gaussian(hsize, sigma):
187
+ hsize = [hsize, hsize]
188
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
189
+ std = sigma
190
+ [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
191
+ arg = -(x * x + y * y) / (2 * std * std)
192
+ h = np.exp(arg)
193
+ h[h < scipy.finfo(float).eps * h.max()] = 0
194
+ sumh = h.sum()
195
+ if sumh != 0:
196
+ h = h / sumh
197
+ return h
198
+
199
+
200
+ def fspecial_laplacian(alpha):
201
+ alpha = max([0, min([alpha, 1])])
202
+ h1 = alpha / (alpha + 1)
203
+ h2 = (1 - alpha) / (alpha + 1)
204
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
205
+ h = np.array(h)
206
+ return h
207
+
208
+
209
+ def fspecial(filter_type, *args, **kwargs):
210
+ '''
211
+ python code from:
212
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
213
+ '''
214
+ if filter_type == 'gaussian':
215
+ return fspecial_gaussian(*args, **kwargs)
216
+ if filter_type == 'laplacian':
217
+ return fspecial_laplacian(*args, **kwargs)
218
+
219
+
220
+ """
221
+ # --------------------------------------------
222
+ # degradation models
223
+ # --------------------------------------------
224
+ """
225
+
226
+
227
+ def bicubic_degradation(x, sf=3):
228
+ '''
229
+ Args:
230
+ x: HxWxC image, [0, 1]
231
+ sf: down-scale factor
232
+ Return:
233
+ bicubicly downsampled LR image
234
+ '''
235
+ x = util.imresize_np(x, scale=1 / sf)
236
+ return x
237
+
238
+
239
+ def srmd_degradation(x, k, sf=3):
240
+ ''' blur + bicubic downsampling
241
+ Args:
242
+ x: HxWxC image, [0, 1]
243
+ k: hxw, double
244
+ sf: down-scale factor
245
+ Return:
246
+ downsampled LR image
247
+ Reference:
248
+ @inproceedings{zhang2018learning,
249
+ title={Learning a single convolutional super-resolution network for multiple degradations},
250
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
251
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
252
+ pages={3262--3271},
253
+ year={2018}
254
+ }
255
+ '''
256
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror'
257
+ x = bicubic_degradation(x, sf=sf)
258
+ return x
259
+
260
+
261
+ def dpsr_degradation(x, k, sf=3):
262
+ ''' bicubic downsampling + blur
263
+ Args:
264
+ x: HxWxC image, [0, 1]
265
+ k: hxw, double
266
+ sf: down-scale factor
267
+ Return:
268
+ downsampled LR image
269
+ Reference:
270
+ @inproceedings{zhang2019deep,
271
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
272
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
273
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
274
+ pages={1671--1681},
275
+ year={2019}
276
+ }
277
+ '''
278
+ x = bicubic_degradation(x, sf=sf)
279
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
280
+ return x
281
+
282
+
283
+ def classical_degradation(x, k, sf=3):
284
+ ''' blur + downsampling
285
+ Args:
286
+ x: HxWxC image, [0, 1]/[0, 255]
287
+ k: hxw, double
288
+ sf: down-scale factor
289
+ Return:
290
+ downsampled LR image
291
+ '''
292
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
293
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
294
+ st = 0
295
+ return x[st::sf, st::sf, ...]
296
+
297
+
298
+ def add_sharpening(img, weight=0.5, radius=50, threshold=10):
299
+ """USM sharpening. borrowed from real-ESRGAN
300
+ Input image: I; Blurry image: B.
301
+ 1. K = I + weight * (I - B)
302
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
303
+ 3. Blur mask:
304
+ 4. Out = Mask * K + (1 - Mask) * I
305
+ Args:
306
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
307
+ weight (float): Sharp weight. Default: 1.
308
+ radius (float): Kernel size of Gaussian blur. Default: 50.
309
+ threshold (int):
310
+ """
311
+ if radius % 2 == 0:
312
+ radius += 1
313
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
314
+ residual = img - blur
315
+ mask = np.abs(residual) * 255 > threshold
316
+ mask = mask.astype('float32')
317
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
318
+
319
+ K = img + weight * residual
320
+ K = np.clip(K, 0, 1)
321
+ return soft_mask * K + (1 - soft_mask) * img
322
+
323
+
324
+ def add_blur(img, sf=4):
325
+ wd2 = 4.0 + sf
326
+ wd = 2.0 + 0.2 * sf
327
+
328
+ wd2 = wd2/4
329
+ wd = wd/4
330
+
331
+ if random.random() < 0.5:
332
+ l1 = wd2 * random.random()
333
+ l2 = wd2 * random.random()
334
+ k = anisotropic_Gaussian(ksize=random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2)
335
+ else:
336
+ k = fspecial('gaussian', random.randint(2, 4) + 3, wd * random.random())
337
+ img = ndimage.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
338
+
339
+ return img
340
+
341
+
342
+ def add_resize(img, sf=4):
343
+ rnum = np.random.rand()
344
+ if rnum > 0.8: # up
345
+ sf1 = random.uniform(1, 2)
346
+ elif rnum < 0.7: # down
347
+ sf1 = random.uniform(0.5 / sf, 1)
348
+ else:
349
+ sf1 = 1.0
350
+ img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3]))
351
+ img = np.clip(img, 0.0, 1.0)
352
+
353
+ return img
354
+
355
+
356
+ # def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
357
+ # noise_level = random.randint(noise_level1, noise_level2)
358
+ # rnum = np.random.rand()
359
+ # if rnum > 0.6: # add color Gaussian noise
360
+ # img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
361
+ # elif rnum < 0.4: # add grayscale Gaussian noise
362
+ # img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
363
+ # else: # add noise
364
+ # L = noise_level2 / 255.
365
+ # D = np.diag(np.random.rand(3))
366
+ # U = orth(np.random.rand(3, 3))
367
+ # conv = np.dot(np.dot(np.transpose(U), D), U)
368
+ # img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
369
+ # img = np.clip(img, 0.0, 1.0)
370
+ # return img
371
+
372
+ def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
373
+ noise_level = random.randint(noise_level1, noise_level2)
374
+ rnum = np.random.rand()
375
+ if rnum > 0.6: # add color Gaussian noise
376
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
377
+ elif rnum < 0.4: # add grayscale Gaussian noise
378
+ img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
379
+ else: # add noise
380
+ L = noise_level2 / 255.
381
+ D = np.diag(np.random.rand(3))
382
+ U = orth(np.random.rand(3, 3))
383
+ conv = np.dot(np.dot(np.transpose(U), D), U)
384
+ img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
385
+ img = np.clip(img, 0.0, 1.0)
386
+ return img
387
+
388
+
389
+ def add_speckle_noise(img, noise_level1=2, noise_level2=25):
390
+ noise_level = random.randint(noise_level1, noise_level2)
391
+ img = np.clip(img, 0.0, 1.0)
392
+ rnum = random.random()
393
+ if rnum > 0.6:
394
+ img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
395
+ elif rnum < 0.4:
396
+ img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
397
+ else:
398
+ L = noise_level2 / 255.
399
+ D = np.diag(np.random.rand(3))
400
+ U = orth(np.random.rand(3, 3))
401
+ conv = np.dot(np.dot(np.transpose(U), D), U)
402
+ img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
403
+ img = np.clip(img, 0.0, 1.0)
404
+ return img
405
+
406
+
407
+ def add_Poisson_noise(img):
408
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.
409
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
410
+ if random.random() < 0.5:
411
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
412
+ else:
413
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
414
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
415
+ noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
416
+ img += noise_gray[:, :, np.newaxis]
417
+ img = np.clip(img, 0.0, 1.0)
418
+ return img
419
+
420
+
421
+ def add_JPEG_noise(img):
422
+ quality_factor = random.randint(80, 95)
423
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
424
+ result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
425
+ img = cv2.imdecode(encimg, 1)
426
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
427
+ return img
428
+
429
+
430
+ def random_crop(lq, hq, sf=4, lq_patchsize=64):
431
+ h, w = lq.shape[:2]
432
+ rnd_h = random.randint(0, h - lq_patchsize)
433
+ rnd_w = random.randint(0, w - lq_patchsize)
434
+ lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
435
+
436
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
437
+ hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
438
+ return lq, hq
439
+
440
+
441
+ def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
442
+ """
443
+ This is the degradation model of BSRGAN from the paper
444
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
445
+ ----------
446
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
447
+ sf: scale factor
448
+ isp_model: camera ISP model
449
+ Returns
450
+ -------
451
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
452
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
453
+ """
454
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
455
+ sf_ori = sf
456
+
457
+ h1, w1 = img.shape[:2]
458
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
459
+ h, w = img.shape[:2]
460
+
461
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
462
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
463
+
464
+ hq = img.copy()
465
+
466
+ if sf == 4 and random.random() < scale2_prob: # downsample1
467
+ if np.random.rand() < 0.5:
468
+ img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
469
+ interpolation=random.choice([1, 2, 3]))
470
+ else:
471
+ img = util.imresize_np(img, 1 / 2, True)
472
+ img = np.clip(img, 0.0, 1.0)
473
+ sf = 2
474
+
475
+ shuffle_order = random.sample(range(7), 7)
476
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
477
+ if idx1 > idx2: # keep downsample3 last
478
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
479
+
480
+ for i in shuffle_order:
481
+
482
+ if i == 0:
483
+ img = add_blur(img, sf=sf)
484
+
485
+ elif i == 1:
486
+ img = add_blur(img, sf=sf)
487
+
488
+ elif i == 2:
489
+ a, b = img.shape[1], img.shape[0]
490
+ # downsample2
491
+ if random.random() < 0.75:
492
+ sf1 = random.uniform(1, 2 * sf)
493
+ img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
494
+ interpolation=random.choice([1, 2, 3]))
495
+ else:
496
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
497
+ k_shifted = shift_pixel(k, sf)
498
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
499
+ img = ndimage.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
500
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
501
+ img = np.clip(img, 0.0, 1.0)
502
+
503
+ elif i == 3:
504
+ # downsample3
505
+ img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
506
+ img = np.clip(img, 0.0, 1.0)
507
+
508
+ elif i == 4:
509
+ # add Gaussian noise
510
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=8)
511
+
512
+ elif i == 5:
513
+ # add JPEG noise
514
+ if random.random() < jpeg_prob:
515
+ img = add_JPEG_noise(img)
516
+
517
+ elif i == 6:
518
+ # add processed camera sensor noise
519
+ if random.random() < isp_prob and isp_model is not None:
520
+ with torch.no_grad():
521
+ img, hq = isp_model.forward(img.copy(), hq)
522
+
523
+ # add final JPEG compression noise
524
+ img = add_JPEG_noise(img)
525
+
526
+ # random crop
527
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
528
+
529
+ return img, hq
530
+
531
+
532
+ # todo no isp_model?
533
+ def degradation_bsrgan_variant(image, sf=4, isp_model=None, up=False):
534
+ """
535
+ This is the degradation model of BSRGAN from the paper
536
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
537
+ ----------
538
+ sf: scale factor
539
+ isp_model: camera ISP model
540
+ Returns
541
+ -------
542
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
543
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
544
+ """
545
+ image = util.uint2single(image)
546
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
547
+ sf_ori = sf
548
+
549
+ h1, w1 = image.shape[:2]
550
+ image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
551
+ h, w = image.shape[:2]
552
+
553
+ hq = image.copy()
554
+
555
+ if sf == 4 and random.random() < scale2_prob: # downsample1
556
+ if np.random.rand() < 0.5:
557
+ image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
558
+ interpolation=random.choice([1, 2, 3]))
559
+ else:
560
+ image = util.imresize_np(image, 1 / 2, True)
561
+ image = np.clip(image, 0.0, 1.0)
562
+ sf = 2
563
+
564
+ shuffle_order = random.sample(range(7), 7)
565
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
566
+ if idx1 > idx2: # keep downsample3 last
567
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
568
+
569
+ for i in shuffle_order:
570
+
571
+ if i == 0:
572
+ image = add_blur(image, sf=sf)
573
+
574
+ # elif i == 1:
575
+ # image = add_blur(image, sf=sf)
576
+
577
+ if i == 0:
578
+ pass
579
+
580
+ elif i == 2:
581
+ a, b = image.shape[1], image.shape[0]
582
+ # downsample2
583
+ if random.random() < 0.8:
584
+ sf1 = random.uniform(1, 2 * sf)
585
+ image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])),
586
+ interpolation=random.choice([1, 2, 3]))
587
+ else:
588
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
589
+ k_shifted = shift_pixel(k, sf)
590
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
591
+ image = ndimage.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror')
592
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
593
+
594
+ image = np.clip(image, 0.0, 1.0)
595
+
596
+ elif i == 3:
597
+ # downsample3
598
+ image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
599
+ image = np.clip(image, 0.0, 1.0)
600
+
601
+ elif i == 4:
602
+ # add Gaussian noise
603
+ image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2)
604
+
605
+ elif i == 5:
606
+ # add JPEG noise
607
+ if random.random() < jpeg_prob:
608
+ image = add_JPEG_noise(image)
609
+ #
610
+ # elif i == 6:
611
+ # # add processed camera sensor noise
612
+ # if random.random() < isp_prob and isp_model is not None:
613
+ # with torch.no_grad():
614
+ # img, hq = isp_model.forward(img.copy(), hq)
615
+
616
+ # add final JPEG compression noise
617
+ image = add_JPEG_noise(image)
618
+ image = util.single2uint(image)
619
+ if up:
620
+ image = cv2.resize(image, (w1, h1), interpolation=cv2.INTER_CUBIC) # todo: random, as above? want to condition on it then
621
+ example = {"image": image}
622
+ return example
623
+
624
+
625
+
626
+
627
+ if __name__ == '__main__':
628
+ print("hey")
629
+ img = util.imread_uint('utils/test.png', 3)
630
+ img = img[:448, :448]
631
+ h = img.shape[0] // 4
632
+ print("resizing to", h)
633
+ sf = 4
634
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
635
+ for i in range(20):
636
+ print(i)
637
+ img_hq = img
638
+ img_lq = deg_fn(img)["image"]
639
+ img_hq, img_lq = util.uint2single(img_hq), util.uint2single(img_lq)
640
+ print(img_lq)
641
+ img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img_hq)["image"]
642
+ print(img_lq.shape)
643
+ print("bicubic", img_lq_bicubic.shape)
644
+ print(img_hq.shape)
645
+ lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
646
+ interpolation=0)
647
+ lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic),
648
+ (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
649
+ interpolation=0)
650
+ img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1)
651
+ util.imsave(img_concat, str(i) + '.png')
comfy/ldm/modules/image_degradation/utils/test.png ADDED
comfy/ldm/modules/image_degradation/utils_image.py ADDED
@@ -0,0 +1,916 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import math
3
+ import random
4
+ import numpy as np
5
+ import torch
6
+ import cv2
7
+ from torchvision.utils import make_grid
8
+ from datetime import datetime
9
+ #import matplotlib.pyplot as plt # TODO: check with Dominik, also bsrgan.py vs bsrgan_light.py
10
+
11
+
12
+ os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
13
+
14
+
15
+ '''
16
+ # --------------------------------------------
17
+ # Kai Zhang (github: https://github.com/cszn)
18
+ # 03/Mar/2019
19
+ # --------------------------------------------
20
+ # https://github.com/twhui/SRGAN-pyTorch
21
+ # https://github.com/xinntao/BasicSR
22
+ # --------------------------------------------
23
+ '''
24
+
25
+
26
+ IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
27
+
28
+
29
+ def is_image_file(filename):
30
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
31
+
32
+
33
+ def get_timestamp():
34
+ return datetime.now().strftime('%y%m%d-%H%M%S')
35
+
36
+
37
+ def imshow(x, title=None, cbar=False, figsize=None):
38
+ plt.figure(figsize=figsize)
39
+ plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
40
+ if title:
41
+ plt.title(title)
42
+ if cbar:
43
+ plt.colorbar()
44
+ plt.show()
45
+
46
+
47
+ def surf(Z, cmap='rainbow', figsize=None):
48
+ plt.figure(figsize=figsize)
49
+ ax3 = plt.axes(projection='3d')
50
+
51
+ w, h = Z.shape[:2]
52
+ xx = np.arange(0,w,1)
53
+ yy = np.arange(0,h,1)
54
+ X, Y = np.meshgrid(xx, yy)
55
+ ax3.plot_surface(X,Y,Z,cmap=cmap)
56
+ #ax3.contour(X,Y,Z, zdim='z',offset=-2,cmap=cmap)
57
+ plt.show()
58
+
59
+
60
+ '''
61
+ # --------------------------------------------
62
+ # get image pathes
63
+ # --------------------------------------------
64
+ '''
65
+
66
+
67
+ def get_image_paths(dataroot):
68
+ paths = None # return None if dataroot is None
69
+ if dataroot is not None:
70
+ paths = sorted(_get_paths_from_images(dataroot))
71
+ return paths
72
+
73
+
74
+ def _get_paths_from_images(path):
75
+ assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
76
+ images = []
77
+ for dirpath, _, fnames in sorted(os.walk(path)):
78
+ for fname in sorted(fnames):
79
+ if is_image_file(fname):
80
+ img_path = os.path.join(dirpath, fname)
81
+ images.append(img_path)
82
+ assert images, '{:s} has no valid image file'.format(path)
83
+ return images
84
+
85
+
86
+ '''
87
+ # --------------------------------------------
88
+ # split large images into small images
89
+ # --------------------------------------------
90
+ '''
91
+
92
+
93
+ def patches_from_image(img, p_size=512, p_overlap=64, p_max=800):
94
+ w, h = img.shape[:2]
95
+ patches = []
96
+ if w > p_max and h > p_max:
97
+ w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int))
98
+ h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int))
99
+ w1.append(w-p_size)
100
+ h1.append(h-p_size)
101
+ # print(w1)
102
+ # print(h1)
103
+ for i in w1:
104
+ for j in h1:
105
+ patches.append(img[i:i+p_size, j:j+p_size,:])
106
+ else:
107
+ patches.append(img)
108
+
109
+ return patches
110
+
111
+
112
+ def imssave(imgs, img_path):
113
+ """
114
+ imgs: list, N images of size WxHxC
115
+ """
116
+ img_name, ext = os.path.splitext(os.path.basename(img_path))
117
+
118
+ for i, img in enumerate(imgs):
119
+ if img.ndim == 3:
120
+ img = img[:, :, [2, 1, 0]]
121
+ new_path = os.path.join(os.path.dirname(img_path), img_name+str('_s{:04d}'.format(i))+'.png')
122
+ cv2.imwrite(new_path, img)
123
+
124
+
125
+ def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=800, p_overlap=96, p_max=1000):
126
+ """
127
+ split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size),
128
+ and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max)
129
+ will be splitted.
130
+ Args:
131
+ original_dataroot:
132
+ taget_dataroot:
133
+ p_size: size of small images
134
+ p_overlap: patch size in training is a good choice
135
+ p_max: images with smaller size than (p_max)x(p_max) keep unchanged.
136
+ """
137
+ paths = get_image_paths(original_dataroot)
138
+ for img_path in paths:
139
+ # img_name, ext = os.path.splitext(os.path.basename(img_path))
140
+ img = imread_uint(img_path, n_channels=n_channels)
141
+ patches = patches_from_image(img, p_size, p_overlap, p_max)
142
+ imssave(patches, os.path.join(taget_dataroot,os.path.basename(img_path)))
143
+ #if original_dataroot == taget_dataroot:
144
+ #del img_path
145
+
146
+ '''
147
+ # --------------------------------------------
148
+ # makedir
149
+ # --------------------------------------------
150
+ '''
151
+
152
+
153
+ def mkdir(path):
154
+ if not os.path.exists(path):
155
+ os.makedirs(path)
156
+
157
+
158
+ def mkdirs(paths):
159
+ if isinstance(paths, str):
160
+ mkdir(paths)
161
+ else:
162
+ for path in paths:
163
+ mkdir(path)
164
+
165
+
166
+ def mkdir_and_rename(path):
167
+ if os.path.exists(path):
168
+ new_name = path + '_archived_' + get_timestamp()
169
+ print('Path already exists. Rename it to [{:s}]'.format(new_name))
170
+ os.rename(path, new_name)
171
+ os.makedirs(path)
172
+
173
+
174
+ '''
175
+ # --------------------------------------------
176
+ # read image from path
177
+ # opencv is fast, but read BGR numpy image
178
+ # --------------------------------------------
179
+ '''
180
+
181
+
182
+ # --------------------------------------------
183
+ # get uint8 image of size HxWxn_channles (RGB)
184
+ # --------------------------------------------
185
+ def imread_uint(path, n_channels=3):
186
+ # input: path
187
+ # output: HxWx3(RGB or GGG), or HxWx1 (G)
188
+ if n_channels == 1:
189
+ img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE
190
+ img = np.expand_dims(img, axis=2) # HxWx1
191
+ elif n_channels == 3:
192
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G
193
+ if img.ndim == 2:
194
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG
195
+ else:
196
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB
197
+ return img
198
+
199
+
200
+ # --------------------------------------------
201
+ # matlab's imwrite
202
+ # --------------------------------------------
203
+ def imsave(img, img_path):
204
+ img = np.squeeze(img)
205
+ if img.ndim == 3:
206
+ img = img[:, :, [2, 1, 0]]
207
+ cv2.imwrite(img_path, img)
208
+
209
+ def imwrite(img, img_path):
210
+ img = np.squeeze(img)
211
+ if img.ndim == 3:
212
+ img = img[:, :, [2, 1, 0]]
213
+ cv2.imwrite(img_path, img)
214
+
215
+
216
+
217
+ # --------------------------------------------
218
+ # get single image of size HxWxn_channles (BGR)
219
+ # --------------------------------------------
220
+ def read_img(path):
221
+ # read image by cv2
222
+ # return: Numpy float32, HWC, BGR, [0,1]
223
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE
224
+ img = img.astype(np.float32) / 255.
225
+ if img.ndim == 2:
226
+ img = np.expand_dims(img, axis=2)
227
+ # some images have 4 channels
228
+ if img.shape[2] > 3:
229
+ img = img[:, :, :3]
230
+ return img
231
+
232
+
233
+ '''
234
+ # --------------------------------------------
235
+ # image format conversion
236
+ # --------------------------------------------
237
+ # numpy(single) <---> numpy(unit)
238
+ # numpy(single) <---> tensor
239
+ # numpy(unit) <---> tensor
240
+ # --------------------------------------------
241
+ '''
242
+
243
+
244
+ # --------------------------------------------
245
+ # numpy(single) [0, 1] <---> numpy(unit)
246
+ # --------------------------------------------
247
+
248
+
249
+ def uint2single(img):
250
+
251
+ return np.float32(img/255.)
252
+
253
+
254
+ def single2uint(img):
255
+
256
+ return np.uint8((img.clip(0, 1)*255.).round())
257
+
258
+
259
+ def uint162single(img):
260
+
261
+ return np.float32(img/65535.)
262
+
263
+
264
+ def single2uint16(img):
265
+
266
+ return np.uint16((img.clip(0, 1)*65535.).round())
267
+
268
+
269
+ # --------------------------------------------
270
+ # numpy(unit) (HxWxC or HxW) <---> tensor
271
+ # --------------------------------------------
272
+
273
+
274
+ # convert uint to 4-dimensional torch tensor
275
+ def uint2tensor4(img):
276
+ if img.ndim == 2:
277
+ img = np.expand_dims(img, axis=2)
278
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
279
+
280
+
281
+ # convert uint to 3-dimensional torch tensor
282
+ def uint2tensor3(img):
283
+ if img.ndim == 2:
284
+ img = np.expand_dims(img, axis=2)
285
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
286
+
287
+
288
+ # convert 2/3/4-dimensional torch tensor to uint
289
+ def tensor2uint(img):
290
+ img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
291
+ if img.ndim == 3:
292
+ img = np.transpose(img, (1, 2, 0))
293
+ return np.uint8((img*255.0).round())
294
+
295
+
296
+ # --------------------------------------------
297
+ # numpy(single) (HxWxC) <---> tensor
298
+ # --------------------------------------------
299
+
300
+
301
+ # convert single (HxWxC) to 3-dimensional torch tensor
302
+ def single2tensor3(img):
303
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
304
+
305
+
306
+ # convert single (HxWxC) to 4-dimensional torch tensor
307
+ def single2tensor4(img):
308
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
309
+
310
+
311
+ # convert torch tensor to single
312
+ def tensor2single(img):
313
+ img = img.data.squeeze().float().cpu().numpy()
314
+ if img.ndim == 3:
315
+ img = np.transpose(img, (1, 2, 0))
316
+
317
+ return img
318
+
319
+ # convert torch tensor to single
320
+ def tensor2single3(img):
321
+ img = img.data.squeeze().float().cpu().numpy()
322
+ if img.ndim == 3:
323
+ img = np.transpose(img, (1, 2, 0))
324
+ elif img.ndim == 2:
325
+ img = np.expand_dims(img, axis=2)
326
+ return img
327
+
328
+
329
+ def single2tensor5(img):
330
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0)
331
+
332
+
333
+ def single32tensor5(img):
334
+ return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0)
335
+
336
+
337
+ def single42tensor4(img):
338
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float()
339
+
340
+
341
+ # from skimage.io import imread, imsave
342
+ def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
343
+ '''
344
+ Converts a torch Tensor into an image Numpy array of BGR channel order
345
+ Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
346
+ Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
347
+ '''
348
+ tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp
349
+ tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1]
350
+ n_dim = tensor.dim()
351
+ if n_dim == 4:
352
+ n_img = len(tensor)
353
+ img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
354
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
355
+ elif n_dim == 3:
356
+ img_np = tensor.numpy()
357
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
358
+ elif n_dim == 2:
359
+ img_np = tensor.numpy()
360
+ else:
361
+ raise TypeError(
362
+ 'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
363
+ if out_type == np.uint8:
364
+ img_np = (img_np * 255.0).round()
365
+ # Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
366
+ return img_np.astype(out_type)
367
+
368
+
369
+ '''
370
+ # --------------------------------------------
371
+ # Augmentation, flipe and/or rotate
372
+ # --------------------------------------------
373
+ # The following two are enough.
374
+ # (1) augmet_img: numpy image of WxHxC or WxH
375
+ # (2) augment_img_tensor4: tensor image 1xCxWxH
376
+ # --------------------------------------------
377
+ '''
378
+
379
+
380
+ def augment_img(img, mode=0):
381
+ '''Kai Zhang (github: https://github.com/cszn)
382
+ '''
383
+ if mode == 0:
384
+ return img
385
+ elif mode == 1:
386
+ return np.flipud(np.rot90(img))
387
+ elif mode == 2:
388
+ return np.flipud(img)
389
+ elif mode == 3:
390
+ return np.rot90(img, k=3)
391
+ elif mode == 4:
392
+ return np.flipud(np.rot90(img, k=2))
393
+ elif mode == 5:
394
+ return np.rot90(img)
395
+ elif mode == 6:
396
+ return np.rot90(img, k=2)
397
+ elif mode == 7:
398
+ return np.flipud(np.rot90(img, k=3))
399
+
400
+
401
+ def augment_img_tensor4(img, mode=0):
402
+ '''Kai Zhang (github: https://github.com/cszn)
403
+ '''
404
+ if mode == 0:
405
+ return img
406
+ elif mode == 1:
407
+ return img.rot90(1, [2, 3]).flip([2])
408
+ elif mode == 2:
409
+ return img.flip([2])
410
+ elif mode == 3:
411
+ return img.rot90(3, [2, 3])
412
+ elif mode == 4:
413
+ return img.rot90(2, [2, 3]).flip([2])
414
+ elif mode == 5:
415
+ return img.rot90(1, [2, 3])
416
+ elif mode == 6:
417
+ return img.rot90(2, [2, 3])
418
+ elif mode == 7:
419
+ return img.rot90(3, [2, 3]).flip([2])
420
+
421
+
422
+ def augment_img_tensor(img, mode=0):
423
+ '''Kai Zhang (github: https://github.com/cszn)
424
+ '''
425
+ img_size = img.size()
426
+ img_np = img.data.cpu().numpy()
427
+ if len(img_size) == 3:
428
+ img_np = np.transpose(img_np, (1, 2, 0))
429
+ elif len(img_size) == 4:
430
+ img_np = np.transpose(img_np, (2, 3, 1, 0))
431
+ img_np = augment_img(img_np, mode=mode)
432
+ img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
433
+ if len(img_size) == 3:
434
+ img_tensor = img_tensor.permute(2, 0, 1)
435
+ elif len(img_size) == 4:
436
+ img_tensor = img_tensor.permute(3, 2, 0, 1)
437
+
438
+ return img_tensor.type_as(img)
439
+
440
+
441
+ def augment_img_np3(img, mode=0):
442
+ if mode == 0:
443
+ return img
444
+ elif mode == 1:
445
+ return img.transpose(1, 0, 2)
446
+ elif mode == 2:
447
+ return img[::-1, :, :]
448
+ elif mode == 3:
449
+ img = img[::-1, :, :]
450
+ img = img.transpose(1, 0, 2)
451
+ return img
452
+ elif mode == 4:
453
+ return img[:, ::-1, :]
454
+ elif mode == 5:
455
+ img = img[:, ::-1, :]
456
+ img = img.transpose(1, 0, 2)
457
+ return img
458
+ elif mode == 6:
459
+ img = img[:, ::-1, :]
460
+ img = img[::-1, :, :]
461
+ return img
462
+ elif mode == 7:
463
+ img = img[:, ::-1, :]
464
+ img = img[::-1, :, :]
465
+ img = img.transpose(1, 0, 2)
466
+ return img
467
+
468
+
469
+ def augment_imgs(img_list, hflip=True, rot=True):
470
+ # horizontal flip OR rotate
471
+ hflip = hflip and random.random() < 0.5
472
+ vflip = rot and random.random() < 0.5
473
+ rot90 = rot and random.random() < 0.5
474
+
475
+ def _augment(img):
476
+ if hflip:
477
+ img = img[:, ::-1, :]
478
+ if vflip:
479
+ img = img[::-1, :, :]
480
+ if rot90:
481
+ img = img.transpose(1, 0, 2)
482
+ return img
483
+
484
+ return [_augment(img) for img in img_list]
485
+
486
+
487
+ '''
488
+ # --------------------------------------------
489
+ # modcrop and shave
490
+ # --------------------------------------------
491
+ '''
492
+
493
+
494
+ def modcrop(img_in, scale):
495
+ # img_in: Numpy, HWC or HW
496
+ img = np.copy(img_in)
497
+ if img.ndim == 2:
498
+ H, W = img.shape
499
+ H_r, W_r = H % scale, W % scale
500
+ img = img[:H - H_r, :W - W_r]
501
+ elif img.ndim == 3:
502
+ H, W, C = img.shape
503
+ H_r, W_r = H % scale, W % scale
504
+ img = img[:H - H_r, :W - W_r, :]
505
+ else:
506
+ raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
507
+ return img
508
+
509
+
510
+ def shave(img_in, border=0):
511
+ # img_in: Numpy, HWC or HW
512
+ img = np.copy(img_in)
513
+ h, w = img.shape[:2]
514
+ img = img[border:h-border, border:w-border]
515
+ return img
516
+
517
+
518
+ '''
519
+ # --------------------------------------------
520
+ # image processing process on numpy image
521
+ # channel_convert(in_c, tar_type, img_list):
522
+ # rgb2ycbcr(img, only_y=True):
523
+ # bgr2ycbcr(img, only_y=True):
524
+ # ycbcr2rgb(img):
525
+ # --------------------------------------------
526
+ '''
527
+
528
+
529
+ def rgb2ycbcr(img, only_y=True):
530
+ '''same as matlab rgb2ycbcr
531
+ only_y: only return Y channel
532
+ Input:
533
+ uint8, [0, 255]
534
+ float, [0, 1]
535
+ '''
536
+ in_img_type = img.dtype
537
+ img.astype(np.float32)
538
+ if in_img_type != np.uint8:
539
+ img *= 255.
540
+ # convert
541
+ if only_y:
542
+ rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
543
+ else:
544
+ rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
545
+ [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
546
+ if in_img_type == np.uint8:
547
+ rlt = rlt.round()
548
+ else:
549
+ rlt /= 255.
550
+ return rlt.astype(in_img_type)
551
+
552
+
553
+ def ycbcr2rgb(img):
554
+ '''same as matlab ycbcr2rgb
555
+ Input:
556
+ uint8, [0, 255]
557
+ float, [0, 1]
558
+ '''
559
+ in_img_type = img.dtype
560
+ img.astype(np.float32)
561
+ if in_img_type != np.uint8:
562
+ img *= 255.
563
+ # convert
564
+ rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
565
+ [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
566
+ if in_img_type == np.uint8:
567
+ rlt = rlt.round()
568
+ else:
569
+ rlt /= 255.
570
+ return rlt.astype(in_img_type)
571
+
572
+
573
+ def bgr2ycbcr(img, only_y=True):
574
+ '''bgr version of rgb2ycbcr
575
+ only_y: only return Y channel
576
+ Input:
577
+ uint8, [0, 255]
578
+ float, [0, 1]
579
+ '''
580
+ in_img_type = img.dtype
581
+ img.astype(np.float32)
582
+ if in_img_type != np.uint8:
583
+ img *= 255.
584
+ # convert
585
+ if only_y:
586
+ rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
587
+ else:
588
+ rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
589
+ [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
590
+ if in_img_type == np.uint8:
591
+ rlt = rlt.round()
592
+ else:
593
+ rlt /= 255.
594
+ return rlt.astype(in_img_type)
595
+
596
+
597
+ def channel_convert(in_c, tar_type, img_list):
598
+ # conversion among BGR, gray and y
599
+ if in_c == 3 and tar_type == 'gray': # BGR to gray
600
+ gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
601
+ return [np.expand_dims(img, axis=2) for img in gray_list]
602
+ elif in_c == 3 and tar_type == 'y': # BGR to y
603
+ y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
604
+ return [np.expand_dims(img, axis=2) for img in y_list]
605
+ elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR
606
+ return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
607
+ else:
608
+ return img_list
609
+
610
+
611
+ '''
612
+ # --------------------------------------------
613
+ # metric, PSNR and SSIM
614
+ # --------------------------------------------
615
+ '''
616
+
617
+
618
+ # --------------------------------------------
619
+ # PSNR
620
+ # --------------------------------------------
621
+ def calculate_psnr(img1, img2, border=0):
622
+ # img1 and img2 have range [0, 255]
623
+ #img1 = img1.squeeze()
624
+ #img2 = img2.squeeze()
625
+ if not img1.shape == img2.shape:
626
+ raise ValueError('Input images must have the same dimensions.')
627
+ h, w = img1.shape[:2]
628
+ img1 = img1[border:h-border, border:w-border]
629
+ img2 = img2[border:h-border, border:w-border]
630
+
631
+ img1 = img1.astype(np.float64)
632
+ img2 = img2.astype(np.float64)
633
+ mse = np.mean((img1 - img2)**2)
634
+ if mse == 0:
635
+ return float('inf')
636
+ return 20 * math.log10(255.0 / math.sqrt(mse))
637
+
638
+
639
+ # --------------------------------------------
640
+ # SSIM
641
+ # --------------------------------------------
642
+ def calculate_ssim(img1, img2, border=0):
643
+ '''calculate SSIM
644
+ the same outputs as MATLAB's
645
+ img1, img2: [0, 255]
646
+ '''
647
+ #img1 = img1.squeeze()
648
+ #img2 = img2.squeeze()
649
+ if not img1.shape == img2.shape:
650
+ raise ValueError('Input images must have the same dimensions.')
651
+ h, w = img1.shape[:2]
652
+ img1 = img1[border:h-border, border:w-border]
653
+ img2 = img2[border:h-border, border:w-border]
654
+
655
+ if img1.ndim == 2:
656
+ return ssim(img1, img2)
657
+ elif img1.ndim == 3:
658
+ if img1.shape[2] == 3:
659
+ ssims = []
660
+ for i in range(3):
661
+ ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
662
+ return np.array(ssims).mean()
663
+ elif img1.shape[2] == 1:
664
+ return ssim(np.squeeze(img1), np.squeeze(img2))
665
+ else:
666
+ raise ValueError('Wrong input image dimensions.')
667
+
668
+
669
+ def ssim(img1, img2):
670
+ C1 = (0.01 * 255)**2
671
+ C2 = (0.03 * 255)**2
672
+
673
+ img1 = img1.astype(np.float64)
674
+ img2 = img2.astype(np.float64)
675
+ kernel = cv2.getGaussianKernel(11, 1.5)
676
+ window = np.outer(kernel, kernel.transpose())
677
+
678
+ mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid
679
+ mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
680
+ mu1_sq = mu1**2
681
+ mu2_sq = mu2**2
682
+ mu1_mu2 = mu1 * mu2
683
+ sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
684
+ sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
685
+ sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
686
+
687
+ ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
688
+ (sigma1_sq + sigma2_sq + C2))
689
+ return ssim_map.mean()
690
+
691
+
692
+ '''
693
+ # --------------------------------------------
694
+ # matlab's bicubic imresize (numpy and torch) [0, 1]
695
+ # --------------------------------------------
696
+ '''
697
+
698
+
699
+ # matlab 'imresize' function, now only support 'bicubic'
700
+ def cubic(x):
701
+ absx = torch.abs(x)
702
+ absx2 = absx**2
703
+ absx3 = absx**3
704
+ return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
705
+ (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
706
+
707
+
708
+ def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
709
+ if (scale < 1) and (antialiasing):
710
+ # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
711
+ kernel_width = kernel_width / scale
712
+
713
+ # Output-space coordinates
714
+ x = torch.linspace(1, out_length, out_length)
715
+
716
+ # Input-space coordinates. Calculate the inverse mapping such that 0.5
717
+ # in output space maps to 0.5 in input space, and 0.5+scale in output
718
+ # space maps to 1.5 in input space.
719
+ u = x / scale + 0.5 * (1 - 1 / scale)
720
+
721
+ # What is the left-most pixel that can be involved in the computation?
722
+ left = torch.floor(u - kernel_width / 2)
723
+
724
+ # What is the maximum number of pixels that can be involved in the
725
+ # computation? Note: it's OK to use an extra pixel here; if the
726
+ # corresponding weights are all zero, it will be eliminated at the end
727
+ # of this function.
728
+ P = math.ceil(kernel_width) + 2
729
+
730
+ # The indices of the input pixels involved in computing the k-th output
731
+ # pixel are in row k of the indices matrix.
732
+ indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
733
+ 1, P).expand(out_length, P)
734
+
735
+ # The weights used to compute the k-th output pixel are in row k of the
736
+ # weights matrix.
737
+ distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
738
+ # apply cubic kernel
739
+ if (scale < 1) and (antialiasing):
740
+ weights = scale * cubic(distance_to_center * scale)
741
+ else:
742
+ weights = cubic(distance_to_center)
743
+ # Normalize the weights matrix so that each row sums to 1.
744
+ weights_sum = torch.sum(weights, 1).view(out_length, 1)
745
+ weights = weights / weights_sum.expand(out_length, P)
746
+
747
+ # If a column in weights is all zero, get rid of it. only consider the first and last column.
748
+ weights_zero_tmp = torch.sum((weights == 0), 0)
749
+ if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
750
+ indices = indices.narrow(1, 1, P - 2)
751
+ weights = weights.narrow(1, 1, P - 2)
752
+ if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
753
+ indices = indices.narrow(1, 0, P - 2)
754
+ weights = weights.narrow(1, 0, P - 2)
755
+ weights = weights.contiguous()
756
+ indices = indices.contiguous()
757
+ sym_len_s = -indices.min() + 1
758
+ sym_len_e = indices.max() - in_length
759
+ indices = indices + sym_len_s - 1
760
+ return weights, indices, int(sym_len_s), int(sym_len_e)
761
+
762
+
763
+ # --------------------------------------------
764
+ # imresize for tensor image [0, 1]
765
+ # --------------------------------------------
766
+ def imresize(img, scale, antialiasing=True):
767
+ # Now the scale should be the same for H and W
768
+ # input: img: pytorch tensor, CHW or HW [0,1]
769
+ # output: CHW or HW [0,1] w/o round
770
+ need_squeeze = True if img.dim() == 2 else False
771
+ if need_squeeze:
772
+ img.unsqueeze_(0)
773
+ in_C, in_H, in_W = img.size()
774
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
775
+ kernel_width = 4
776
+ kernel = 'cubic'
777
+
778
+ # Return the desired dimension order for performing the resize. The
779
+ # strategy is to perform the resize first along the dimension with the
780
+ # smallest scale factor.
781
+ # Now we do not support this.
782
+
783
+ # get weights and indices
784
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
785
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
786
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
787
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
788
+ # process H dimension
789
+ # symmetric copying
790
+ img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
791
+ img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
792
+
793
+ sym_patch = img[:, :sym_len_Hs, :]
794
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
795
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
796
+ img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
797
+
798
+ sym_patch = img[:, -sym_len_He:, :]
799
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
800
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
801
+ img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
802
+
803
+ out_1 = torch.FloatTensor(in_C, out_H, in_W)
804
+ kernel_width = weights_H.size(1)
805
+ for i in range(out_H):
806
+ idx = int(indices_H[i][0])
807
+ for j in range(out_C):
808
+ out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
809
+
810
+ # process W dimension
811
+ # symmetric copying
812
+ out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
813
+ out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
814
+
815
+ sym_patch = out_1[:, :, :sym_len_Ws]
816
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
817
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
818
+ out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
819
+
820
+ sym_patch = out_1[:, :, -sym_len_We:]
821
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
822
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
823
+ out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
824
+
825
+ out_2 = torch.FloatTensor(in_C, out_H, out_W)
826
+ kernel_width = weights_W.size(1)
827
+ for i in range(out_W):
828
+ idx = int(indices_W[i][0])
829
+ for j in range(out_C):
830
+ out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])
831
+ if need_squeeze:
832
+ out_2.squeeze_()
833
+ return out_2
834
+
835
+
836
+ # --------------------------------------------
837
+ # imresize for numpy image [0, 1]
838
+ # --------------------------------------------
839
+ def imresize_np(img, scale, antialiasing=True):
840
+ # Now the scale should be the same for H and W
841
+ # input: img: Numpy, HWC or HW [0,1]
842
+ # output: HWC or HW [0,1] w/o round
843
+ img = torch.from_numpy(img)
844
+ need_squeeze = True if img.dim() == 2 else False
845
+ if need_squeeze:
846
+ img.unsqueeze_(2)
847
+
848
+ in_H, in_W, in_C = img.size()
849
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
850
+ kernel_width = 4
851
+ kernel = 'cubic'
852
+
853
+ # Return the desired dimension order for performing the resize. The
854
+ # strategy is to perform the resize first along the dimension with the
855
+ # smallest scale factor.
856
+ # Now we do not support this.
857
+
858
+ # get weights and indices
859
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
860
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
861
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
862
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
863
+ # process H dimension
864
+ # symmetric copying
865
+ img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
866
+ img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
867
+
868
+ sym_patch = img[:sym_len_Hs, :, :]
869
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
870
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
871
+ img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
872
+
873
+ sym_patch = img[-sym_len_He:, :, :]
874
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
875
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
876
+ img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
877
+
878
+ out_1 = torch.FloatTensor(out_H, in_W, in_C)
879
+ kernel_width = weights_H.size(1)
880
+ for i in range(out_H):
881
+ idx = int(indices_H[i][0])
882
+ for j in range(out_C):
883
+ out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
884
+
885
+ # process W dimension
886
+ # symmetric copying
887
+ out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
888
+ out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
889
+
890
+ sym_patch = out_1[:, :sym_len_Ws, :]
891
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
892
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
893
+ out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
894
+
895
+ sym_patch = out_1[:, -sym_len_We:, :]
896
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
897
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
898
+ out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
899
+
900
+ out_2 = torch.FloatTensor(out_H, out_W, in_C)
901
+ kernel_width = weights_W.size(1)
902
+ for i in range(out_W):
903
+ idx = int(indices_W[i][0])
904
+ for j in range(out_C):
905
+ out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
906
+ if need_squeeze:
907
+ out_2.squeeze_()
908
+
909
+ return out_2.numpy()
910
+
911
+
912
+ if __name__ == '__main__':
913
+ print('---')
914
+ # img = imread_uint('test.bmp', 3)
915
+ # img = uint2single(img)
916
+ # img_bicubic = imresize_np(img, 1/4)
comfy/ldm/modules/midas/__init__.py ADDED
File without changes
comfy/ldm/modules/midas/api.py ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # based on https://github.com/isl-org/MiDaS
2
+
3
+ import cv2
4
+ import torch
5
+ import torch.nn as nn
6
+ from torchvision.transforms import Compose
7
+
8
+ from ldm.modules.midas.midas.dpt_depth import DPTDepthModel
9
+ from ldm.modules.midas.midas.midas_net import MidasNet
10
+ from ldm.modules.midas.midas.midas_net_custom import MidasNet_small
11
+ from ldm.modules.midas.midas.transforms import Resize, NormalizeImage, PrepareForNet
12
+
13
+
14
+ ISL_PATHS = {
15
+ "dpt_large": "midas_models/dpt_large-midas-2f21e586.pt",
16
+ "dpt_hybrid": "midas_models/dpt_hybrid-midas-501f0c75.pt",
17
+ "midas_v21": "",
18
+ "midas_v21_small": "",
19
+ }
20
+
21
+
22
+ def disabled_train(self, mode=True):
23
+ """Overwrite model.train with this function to make sure train/eval mode
24
+ does not change anymore."""
25
+ return self
26
+
27
+
28
+ def load_midas_transform(model_type):
29
+ # https://github.com/isl-org/MiDaS/blob/master/run.py
30
+ # load transform only
31
+ if model_type == "dpt_large": # DPT-Large
32
+ net_w, net_h = 384, 384
33
+ resize_mode = "minimal"
34
+ normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
35
+
36
+ elif model_type == "dpt_hybrid": # DPT-Hybrid
37
+ net_w, net_h = 384, 384
38
+ resize_mode = "minimal"
39
+ normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
40
+
41
+ elif model_type == "midas_v21":
42
+ net_w, net_h = 384, 384
43
+ resize_mode = "upper_bound"
44
+ normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
45
+
46
+ elif model_type == "midas_v21_small":
47
+ net_w, net_h = 256, 256
48
+ resize_mode = "upper_bound"
49
+ normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
50
+
51
+ else:
52
+ assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
53
+
54
+ transform = Compose(
55
+ [
56
+ Resize(
57
+ net_w,
58
+ net_h,
59
+ resize_target=None,
60
+ keep_aspect_ratio=True,
61
+ ensure_multiple_of=32,
62
+ resize_method=resize_mode,
63
+ image_interpolation_method=cv2.INTER_CUBIC,
64
+ ),
65
+ normalization,
66
+ PrepareForNet(),
67
+ ]
68
+ )
69
+
70
+ return transform
71
+
72
+
73
+ def load_model(model_type):
74
+ # https://github.com/isl-org/MiDaS/blob/master/run.py
75
+ # load network
76
+ model_path = ISL_PATHS[model_type]
77
+ if model_type == "dpt_large": # DPT-Large
78
+ model = DPTDepthModel(
79
+ path=model_path,
80
+ backbone="vitl16_384",
81
+ non_negative=True,
82
+ )
83
+ net_w, net_h = 384, 384
84
+ resize_mode = "minimal"
85
+ normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
86
+
87
+ elif model_type == "dpt_hybrid": # DPT-Hybrid
88
+ model = DPTDepthModel(
89
+ path=model_path,
90
+ backbone="vitb_rn50_384",
91
+ non_negative=True,
92
+ )
93
+ net_w, net_h = 384, 384
94
+ resize_mode = "minimal"
95
+ normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
96
+
97
+ elif model_type == "midas_v21":
98
+ model = MidasNet(model_path, non_negative=True)
99
+ net_w, net_h = 384, 384
100
+ resize_mode = "upper_bound"
101
+ normalization = NormalizeImage(
102
+ mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
103
+ )
104
+
105
+ elif model_type == "midas_v21_small":
106
+ model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
107
+ non_negative=True, blocks={'expand': True})
108
+ net_w, net_h = 256, 256
109
+ resize_mode = "upper_bound"
110
+ normalization = NormalizeImage(
111
+ mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
112
+ )
113
+
114
+ else:
115
+ print(f"model_type '{model_type}' not implemented, use: --model_type large")
116
+ assert False
117
+
118
+ transform = Compose(
119
+ [
120
+ Resize(
121
+ net_w,
122
+ net_h,
123
+ resize_target=None,
124
+ keep_aspect_ratio=True,
125
+ ensure_multiple_of=32,
126
+ resize_method=resize_mode,
127
+ image_interpolation_method=cv2.INTER_CUBIC,
128
+ ),
129
+ normalization,
130
+ PrepareForNet(),
131
+ ]
132
+ )
133
+
134
+ return model.eval(), transform
135
+
136
+
137
+ class MiDaSInference(nn.Module):
138
+ MODEL_TYPES_TORCH_HUB = [
139
+ "DPT_Large",
140
+ "DPT_Hybrid",
141
+ "MiDaS_small"
142
+ ]
143
+ MODEL_TYPES_ISL = [
144
+ "dpt_large",
145
+ "dpt_hybrid",
146
+ "midas_v21",
147
+ "midas_v21_small",
148
+ ]
149
+
150
+ def __init__(self, model_type):
151
+ super().__init__()
152
+ assert (model_type in self.MODEL_TYPES_ISL)
153
+ model, _ = load_model(model_type)
154
+ self.model = model
155
+ self.model.train = disabled_train
156
+
157
+ def forward(self, x):
158
+ # x in 0..1 as produced by calling self.transform on a 0..1 float64 numpy array
159
+ # NOTE: we expect that the correct transform has been called during dataloading.
160
+ with torch.no_grad():
161
+ prediction = self.model(x)
162
+ prediction = torch.nn.functional.interpolate(
163
+ prediction.unsqueeze(1),
164
+ size=x.shape[2:],
165
+ mode="bicubic",
166
+ align_corners=False,
167
+ )
168
+ assert prediction.shape == (x.shape[0], 1, x.shape[2], x.shape[3])
169
+ return prediction
170
+
comfy/ldm/modules/midas/midas/__init__.py ADDED
File without changes
comfy/ldm/modules/midas/midas/base_model.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+
4
+ class BaseModel(torch.nn.Module):
5
+ def load(self, path):
6
+ """Load model from file.
7
+
8
+ Args:
9
+ path (str): file path
10
+ """
11
+ parameters = torch.load(path, map_location=torch.device('cpu'))
12
+
13
+ if "optimizer" in parameters:
14
+ parameters = parameters["model"]
15
+
16
+ self.load_state_dict(parameters)
comfy/ldm/modules/midas/midas/blocks.py ADDED
@@ -0,0 +1,342 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ from .vit import (
5
+ _make_pretrained_vitb_rn50_384,
6
+ _make_pretrained_vitl16_384,
7
+ _make_pretrained_vitb16_384,
8
+ forward_vit,
9
+ )
10
+
11
+ def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True, hooks=None, use_vit_only=False, use_readout="ignore",):
12
+ if backbone == "vitl16_384":
13
+ pretrained = _make_pretrained_vitl16_384(
14
+ use_pretrained, hooks=hooks, use_readout=use_readout
15
+ )
16
+ scratch = _make_scratch(
17
+ [256, 512, 1024, 1024], features, groups=groups, expand=expand
18
+ ) # ViT-L/16 - 85.0% Top1 (backbone)
19
+ elif backbone == "vitb_rn50_384":
20
+ pretrained = _make_pretrained_vitb_rn50_384(
21
+ use_pretrained,
22
+ hooks=hooks,
23
+ use_vit_only=use_vit_only,
24
+ use_readout=use_readout,
25
+ )
26
+ scratch = _make_scratch(
27
+ [256, 512, 768, 768], features, groups=groups, expand=expand
28
+ ) # ViT-H/16 - 85.0% Top1 (backbone)
29
+ elif backbone == "vitb16_384":
30
+ pretrained = _make_pretrained_vitb16_384(
31
+ use_pretrained, hooks=hooks, use_readout=use_readout
32
+ )
33
+ scratch = _make_scratch(
34
+ [96, 192, 384, 768], features, groups=groups, expand=expand
35
+ ) # ViT-B/16 - 84.6% Top1 (backbone)
36
+ elif backbone == "resnext101_wsl":
37
+ pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
38
+ scratch = _make_scratch([256, 512, 1024, 2048], features, groups=groups, expand=expand) # efficientnet_lite3
39
+ elif backbone == "efficientnet_lite3":
40
+ pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable)
41
+ scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand) # efficientnet_lite3
42
+ else:
43
+ print(f"Backbone '{backbone}' not implemented")
44
+ assert False
45
+
46
+ return pretrained, scratch
47
+
48
+
49
+ def _make_scratch(in_shape, out_shape, groups=1, expand=False):
50
+ scratch = nn.Module()
51
+
52
+ out_shape1 = out_shape
53
+ out_shape2 = out_shape
54
+ out_shape3 = out_shape
55
+ out_shape4 = out_shape
56
+ if expand==True:
57
+ out_shape1 = out_shape
58
+ out_shape2 = out_shape*2
59
+ out_shape3 = out_shape*4
60
+ out_shape4 = out_shape*8
61
+
62
+ scratch.layer1_rn = nn.Conv2d(
63
+ in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
64
+ )
65
+ scratch.layer2_rn = nn.Conv2d(
66
+ in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
67
+ )
68
+ scratch.layer3_rn = nn.Conv2d(
69
+ in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
70
+ )
71
+ scratch.layer4_rn = nn.Conv2d(
72
+ in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
73
+ )
74
+
75
+ return scratch
76
+
77
+
78
+ def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
79
+ efficientnet = torch.hub.load(
80
+ "rwightman/gen-efficientnet-pytorch",
81
+ "tf_efficientnet_lite3",
82
+ pretrained=use_pretrained,
83
+ exportable=exportable
84
+ )
85
+ return _make_efficientnet_backbone(efficientnet)
86
+
87
+
88
+ def _make_efficientnet_backbone(effnet):
89
+ pretrained = nn.Module()
90
+
91
+ pretrained.layer1 = nn.Sequential(
92
+ effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
93
+ )
94
+ pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
95
+ pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
96
+ pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
97
+
98
+ return pretrained
99
+
100
+
101
+ def _make_resnet_backbone(resnet):
102
+ pretrained = nn.Module()
103
+ pretrained.layer1 = nn.Sequential(
104
+ resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
105
+ )
106
+
107
+ pretrained.layer2 = resnet.layer2
108
+ pretrained.layer3 = resnet.layer3
109
+ pretrained.layer4 = resnet.layer4
110
+
111
+ return pretrained
112
+
113
+
114
+ def _make_pretrained_resnext101_wsl(use_pretrained):
115
+ resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
116
+ return _make_resnet_backbone(resnet)
117
+
118
+
119
+
120
+ class Interpolate(nn.Module):
121
+ """Interpolation module.
122
+ """
123
+
124
+ def __init__(self, scale_factor, mode, align_corners=False):
125
+ """Init.
126
+
127
+ Args:
128
+ scale_factor (float): scaling
129
+ mode (str): interpolation mode
130
+ """
131
+ super(Interpolate, self).__init__()
132
+
133
+ self.interp = nn.functional.interpolate
134
+ self.scale_factor = scale_factor
135
+ self.mode = mode
136
+ self.align_corners = align_corners
137
+
138
+ def forward(self, x):
139
+ """Forward pass.
140
+
141
+ Args:
142
+ x (tensor): input
143
+
144
+ Returns:
145
+ tensor: interpolated data
146
+ """
147
+
148
+ x = self.interp(
149
+ x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
150
+ )
151
+
152
+ return x
153
+
154
+
155
+ class ResidualConvUnit(nn.Module):
156
+ """Residual convolution module.
157
+ """
158
+
159
+ def __init__(self, features):
160
+ """Init.
161
+
162
+ Args:
163
+ features (int): number of features
164
+ """
165
+ super().__init__()
166
+
167
+ self.conv1 = nn.Conv2d(
168
+ features, features, kernel_size=3, stride=1, padding=1, bias=True
169
+ )
170
+
171
+ self.conv2 = nn.Conv2d(
172
+ features, features, kernel_size=3, stride=1, padding=1, bias=True
173
+ )
174
+
175
+ self.relu = nn.ReLU(inplace=True)
176
+
177
+ def forward(self, x):
178
+ """Forward pass.
179
+
180
+ Args:
181
+ x (tensor): input
182
+
183
+ Returns:
184
+ tensor: output
185
+ """
186
+ out = self.relu(x)
187
+ out = self.conv1(out)
188
+ out = self.relu(out)
189
+ out = self.conv2(out)
190
+
191
+ return out + x
192
+
193
+
194
+ class FeatureFusionBlock(nn.Module):
195
+ """Feature fusion block.
196
+ """
197
+
198
+ def __init__(self, features):
199
+ """Init.
200
+
201
+ Args:
202
+ features (int): number of features
203
+ """
204
+ super(FeatureFusionBlock, self).__init__()
205
+
206
+ self.resConfUnit1 = ResidualConvUnit(features)
207
+ self.resConfUnit2 = ResidualConvUnit(features)
208
+
209
+ def forward(self, *xs):
210
+ """Forward pass.
211
+
212
+ Returns:
213
+ tensor: output
214
+ """
215
+ output = xs[0]
216
+
217
+ if len(xs) == 2:
218
+ output += self.resConfUnit1(xs[1])
219
+
220
+ output = self.resConfUnit2(output)
221
+
222
+ output = nn.functional.interpolate(
223
+ output, scale_factor=2, mode="bilinear", align_corners=True
224
+ )
225
+
226
+ return output
227
+
228
+
229
+
230
+
231
+ class ResidualConvUnit_custom(nn.Module):
232
+ """Residual convolution module.
233
+ """
234
+
235
+ def __init__(self, features, activation, bn):
236
+ """Init.
237
+
238
+ Args:
239
+ features (int): number of features
240
+ """
241
+ super().__init__()
242
+
243
+ self.bn = bn
244
+
245
+ self.groups=1
246
+
247
+ self.conv1 = nn.Conv2d(
248
+ features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
249
+ )
250
+
251
+ self.conv2 = nn.Conv2d(
252
+ features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
253
+ )
254
+
255
+ if self.bn==True:
256
+ self.bn1 = nn.BatchNorm2d(features)
257
+ self.bn2 = nn.BatchNorm2d(features)
258
+
259
+ self.activation = activation
260
+
261
+ self.skip_add = nn.quantized.FloatFunctional()
262
+
263
+ def forward(self, x):
264
+ """Forward pass.
265
+
266
+ Args:
267
+ x (tensor): input
268
+
269
+ Returns:
270
+ tensor: output
271
+ """
272
+
273
+ out = self.activation(x)
274
+ out = self.conv1(out)
275
+ if self.bn==True:
276
+ out = self.bn1(out)
277
+
278
+ out = self.activation(out)
279
+ out = self.conv2(out)
280
+ if self.bn==True:
281
+ out = self.bn2(out)
282
+
283
+ if self.groups > 1:
284
+ out = self.conv_merge(out)
285
+
286
+ return self.skip_add.add(out, x)
287
+
288
+ # return out + x
289
+
290
+
291
+ class FeatureFusionBlock_custom(nn.Module):
292
+ """Feature fusion block.
293
+ """
294
+
295
+ def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
296
+ """Init.
297
+
298
+ Args:
299
+ features (int): number of features
300
+ """
301
+ super(FeatureFusionBlock_custom, self).__init__()
302
+
303
+ self.deconv = deconv
304
+ self.align_corners = align_corners
305
+
306
+ self.groups=1
307
+
308
+ self.expand = expand
309
+ out_features = features
310
+ if self.expand==True:
311
+ out_features = features//2
312
+
313
+ self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
314
+
315
+ self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
316
+ self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
317
+
318
+ self.skip_add = nn.quantized.FloatFunctional()
319
+
320
+ def forward(self, *xs):
321
+ """Forward pass.
322
+
323
+ Returns:
324
+ tensor: output
325
+ """
326
+ output = xs[0]
327
+
328
+ if len(xs) == 2:
329
+ res = self.resConfUnit1(xs[1])
330
+ output = self.skip_add.add(output, res)
331
+ # output += res
332
+
333
+ output = self.resConfUnit2(output)
334
+
335
+ output = nn.functional.interpolate(
336
+ output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
337
+ )
338
+
339
+ output = self.out_conv(output)
340
+
341
+ return output
342
+
comfy/ldm/modules/midas/midas/dpt_depth.py ADDED
@@ -0,0 +1,109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+ from .base_model import BaseModel
6
+ from .blocks import (
7
+ FeatureFusionBlock,
8
+ FeatureFusionBlock_custom,
9
+ Interpolate,
10
+ _make_encoder,
11
+ forward_vit,
12
+ )
13
+
14
+
15
+ def _make_fusion_block(features, use_bn):
16
+ return FeatureFusionBlock_custom(
17
+ features,
18
+ nn.ReLU(False),
19
+ deconv=False,
20
+ bn=use_bn,
21
+ expand=False,
22
+ align_corners=True,
23
+ )
24
+
25
+
26
+ class DPT(BaseModel):
27
+ def __init__(
28
+ self,
29
+ head,
30
+ features=256,
31
+ backbone="vitb_rn50_384",
32
+ readout="project",
33
+ channels_last=False,
34
+ use_bn=False,
35
+ ):
36
+
37
+ super(DPT, self).__init__()
38
+
39
+ self.channels_last = channels_last
40
+
41
+ hooks = {
42
+ "vitb_rn50_384": [0, 1, 8, 11],
43
+ "vitb16_384": [2, 5, 8, 11],
44
+ "vitl16_384": [5, 11, 17, 23],
45
+ }
46
+
47
+ # Instantiate backbone and reassemble blocks
48
+ self.pretrained, self.scratch = _make_encoder(
49
+ backbone,
50
+ features,
51
+ False, # Set to true of you want to train from scratch, uses ImageNet weights
52
+ groups=1,
53
+ expand=False,
54
+ exportable=False,
55
+ hooks=hooks[backbone],
56
+ use_readout=readout,
57
+ )
58
+
59
+ self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
60
+ self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
61
+ self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
62
+ self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
63
+
64
+ self.scratch.output_conv = head
65
+
66
+
67
+ def forward(self, x):
68
+ if self.channels_last == True:
69
+ x.contiguous(memory_format=torch.channels_last)
70
+
71
+ layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
72
+
73
+ layer_1_rn = self.scratch.layer1_rn(layer_1)
74
+ layer_2_rn = self.scratch.layer2_rn(layer_2)
75
+ layer_3_rn = self.scratch.layer3_rn(layer_3)
76
+ layer_4_rn = self.scratch.layer4_rn(layer_4)
77
+
78
+ path_4 = self.scratch.refinenet4(layer_4_rn)
79
+ path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
80
+ path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
81
+ path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
82
+
83
+ out = self.scratch.output_conv(path_1)
84
+
85
+ return out
86
+
87
+
88
+ class DPTDepthModel(DPT):
89
+ def __init__(self, path=None, non_negative=True, **kwargs):
90
+ features = kwargs["features"] if "features" in kwargs else 256
91
+
92
+ head = nn.Sequential(
93
+ nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
94
+ Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
95
+ nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
96
+ nn.ReLU(True),
97
+ nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
98
+ nn.ReLU(True) if non_negative else nn.Identity(),
99
+ nn.Identity(),
100
+ )
101
+
102
+ super().__init__(head, **kwargs)
103
+
104
+ if path is not None:
105
+ self.load(path)
106
+
107
+ def forward(self, x):
108
+ return super().forward(x).squeeze(dim=1)
109
+