spideyrim commited on
Commit
436faa6
1 Parent(s): fda1a73

Upload 202 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/checkpoint_pickle.py +13 -0
  2. comfy/cldm/cldm.py +308 -0
  3. comfy/cli_args.py +106 -0
  4. comfy/clip_config_bigg.json +23 -0
  5. comfy/clip_vision.py +114 -0
  6. comfy/clip_vision_config_g.json +18 -0
  7. comfy/clip_vision_config_h.json +18 -0
  8. comfy/clip_vision_config_vitl.json +18 -0
  9. comfy/controlnet.py +488 -0
  10. comfy/diffusers_convert.py +261 -0
  11. comfy/diffusers_load.py +36 -0
  12. comfy/extra_samplers/uni_pc.py +883 -0
  13. comfy/gligen.py +341 -0
  14. comfy/k_diffusion/external.py +190 -0
  15. comfy/k_diffusion/sampling.py +739 -0
  16. comfy/k_diffusion/utils.py +313 -0
  17. comfy/latent_formats.py +35 -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 +418 -0
  21. comfy/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  22. comfy/ldm/models/diffusion/dpm_solver/dpm_solver.py +1163 -0
  23. comfy/ldm/models/diffusion/dpm_solver/sampler.py +96 -0
  24. comfy/ldm/models/diffusion/plms.py +245 -0
  25. comfy/ldm/models/diffusion/sampling_util.py +22 -0
  26. comfy/ldm/modules/attention.py +700 -0
  27. comfy/ldm/modules/diffusionmodules/__init__.py +0 -0
  28. comfy/ldm/modules/diffusionmodules/model.py +737 -0
  29. comfy/ldm/modules/diffusionmodules/openaimodel.py +664 -0
  30. comfy/ldm/modules/diffusionmodules/upscaling.py +81 -0
  31. comfy/ldm/modules/diffusionmodules/util.py +278 -0
  32. comfy/ldm/modules/distributions/__init__.py +0 -0
  33. comfy/ldm/modules/distributions/distributions.py +92 -0
  34. comfy/ldm/modules/ema.py +80 -0
  35. comfy/ldm/modules/encoders/__init__.py +0 -0
  36. comfy/ldm/modules/encoders/noise_aug_modules.py +35 -0
  37. comfy/ldm/modules/sub_quadratic_attention.py +250 -0
  38. comfy/ldm/util.py +197 -0
  39. comfy/lora.py +199 -0
  40. comfy/model_base.py +210 -0
  41. comfy/model_detection.py +210 -0
  42. comfy/model_management.py +711 -0
  43. comfy/model_patcher.py +288 -0
  44. comfy/ops.py +46 -0
  45. comfy/options.py +6 -0
  46. comfy/sample.py +97 -0
  47. comfy/samplers.py +744 -0
  48. comfy/sd.py +486 -0
  49. comfy/sd1_clip.py +450 -0
  50. comfy/sd1_clip_config.json +25 -0
comfy/checkpoint_pickle.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pickle
2
+
3
+ load = pickle.load
4
+
5
+ class Empty:
6
+ pass
7
+
8
+ class Unpickler(pickle.Unpickler):
9
+ def find_class(self, module, name):
10
+ #TODO: safe unpickle
11
+ if module.startswith("pytorch_lightning"):
12
+ return Empty
13
+ return super().find_class(module, name)
comfy/cldm/cldm.py ADDED
@@ -0,0 +1,308 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ zero_module,
10
+ timestep_embedding,
11
+ )
12
+
13
+ from ..ldm.modules.attention import SpatialTransformer
14
+ from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
15
+ from ..ldm.util import exists
16
+ import comfy.ops
17
+
18
+ class ControlledUnetModel(UNetModel):
19
+ #implemented in the ldm unet
20
+ pass
21
+
22
+ class ControlNet(nn.Module):
23
+ def __init__(
24
+ self,
25
+ image_size,
26
+ in_channels,
27
+ model_channels,
28
+ hint_channels,
29
+ num_res_blocks,
30
+ attention_resolutions,
31
+ dropout=0,
32
+ channel_mult=(1, 2, 4, 8),
33
+ conv_resample=True,
34
+ dims=2,
35
+ num_classes=None,
36
+ use_checkpoint=False,
37
+ use_fp16=False,
38
+ use_bf16=False,
39
+ num_heads=-1,
40
+ num_head_channels=-1,
41
+ num_heads_upsample=-1,
42
+ use_scale_shift_norm=False,
43
+ resblock_updown=False,
44
+ use_new_attention_order=False,
45
+ use_spatial_transformer=False, # custom transformer support
46
+ transformer_depth=1, # custom transformer support
47
+ context_dim=None, # custom transformer support
48
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
49
+ legacy=True,
50
+ disable_self_attentions=None,
51
+ num_attention_blocks=None,
52
+ disable_middle_self_attn=False,
53
+ use_linear_in_transformer=False,
54
+ adm_in_channels=None,
55
+ transformer_depth_middle=None,
56
+ device=None,
57
+ operations=comfy.ops,
58
+ ):
59
+ super().__init__()
60
+ assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
61
+ if use_spatial_transformer:
62
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
63
+
64
+ if context_dim is not None:
65
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
66
+ # from omegaconf.listconfig import ListConfig
67
+ # if type(context_dim) == ListConfig:
68
+ # context_dim = list(context_dim)
69
+
70
+ if num_heads_upsample == -1:
71
+ num_heads_upsample = num_heads
72
+
73
+ if num_heads == -1:
74
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
75
+
76
+ if num_head_channels == -1:
77
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
78
+
79
+ self.dims = dims
80
+ self.image_size = image_size
81
+ self.in_channels = in_channels
82
+ self.model_channels = model_channels
83
+ if isinstance(transformer_depth, int):
84
+ transformer_depth = len(channel_mult) * [transformer_depth]
85
+ if transformer_depth_middle is None:
86
+ transformer_depth_middle = transformer_depth[-1]
87
+ if isinstance(num_res_blocks, int):
88
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
89
+ else:
90
+ if len(num_res_blocks) != len(channel_mult):
91
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
92
+ "as a list/tuple (per-level) with the same length as channel_mult")
93
+ self.num_res_blocks = num_res_blocks
94
+ if disable_self_attentions is not None:
95
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
96
+ assert len(disable_self_attentions) == len(channel_mult)
97
+ if num_attention_blocks is not None:
98
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
99
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
100
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
101
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
102
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
103
+ f"attention will still not be set.")
104
+
105
+ self.attention_resolutions = attention_resolutions
106
+ self.dropout = dropout
107
+ self.channel_mult = channel_mult
108
+ self.conv_resample = conv_resample
109
+ self.num_classes = num_classes
110
+ self.use_checkpoint = use_checkpoint
111
+ self.dtype = th.float16 if use_fp16 else th.float32
112
+ self.dtype = th.bfloat16 if use_bf16 else self.dtype
113
+ self.num_heads = num_heads
114
+ self.num_head_channels = num_head_channels
115
+ self.num_heads_upsample = num_heads_upsample
116
+ self.predict_codebook_ids = n_embed is not None
117
+
118
+ time_embed_dim = model_channels * 4
119
+ self.time_embed = nn.Sequential(
120
+ operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
121
+ nn.SiLU(),
122
+ operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
123
+ )
124
+
125
+ if self.num_classes is not None:
126
+ if isinstance(self.num_classes, int):
127
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
128
+ elif self.num_classes == "continuous":
129
+ print("setting up linear c_adm embedding layer")
130
+ self.label_emb = nn.Linear(1, time_embed_dim)
131
+ elif self.num_classes == "sequential":
132
+ assert adm_in_channels is not None
133
+ self.label_emb = nn.Sequential(
134
+ nn.Sequential(
135
+ operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
136
+ nn.SiLU(),
137
+ operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
138
+ )
139
+ )
140
+ else:
141
+ raise ValueError()
142
+
143
+ self.input_blocks = nn.ModuleList(
144
+ [
145
+ TimestepEmbedSequential(
146
+ operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
147
+ )
148
+ ]
149
+ )
150
+ self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations)])
151
+
152
+ self.input_hint_block = TimestepEmbedSequential(
153
+ operations.conv_nd(dims, hint_channels, 16, 3, padding=1),
154
+ nn.SiLU(),
155
+ operations.conv_nd(dims, 16, 16, 3, padding=1),
156
+ nn.SiLU(),
157
+ operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2),
158
+ nn.SiLU(),
159
+ operations.conv_nd(dims, 32, 32, 3, padding=1),
160
+ nn.SiLU(),
161
+ operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2),
162
+ nn.SiLU(),
163
+ operations.conv_nd(dims, 96, 96, 3, padding=1),
164
+ nn.SiLU(),
165
+ operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2),
166
+ nn.SiLU(),
167
+ zero_module(operations.conv_nd(dims, 256, model_channels, 3, padding=1))
168
+ )
169
+
170
+ self._feature_size = model_channels
171
+ input_block_chans = [model_channels]
172
+ ch = model_channels
173
+ ds = 1
174
+ for level, mult in enumerate(channel_mult):
175
+ for nr in range(self.num_res_blocks[level]):
176
+ layers = [
177
+ ResBlock(
178
+ ch,
179
+ time_embed_dim,
180
+ dropout,
181
+ out_channels=mult * model_channels,
182
+ dims=dims,
183
+ use_checkpoint=use_checkpoint,
184
+ use_scale_shift_norm=use_scale_shift_norm,
185
+ operations=operations
186
+ )
187
+ ]
188
+ ch = mult * model_channels
189
+ if ds in attention_resolutions:
190
+ if num_head_channels == -1:
191
+ dim_head = ch // num_heads
192
+ else:
193
+ num_heads = ch // num_head_channels
194
+ dim_head = num_head_channels
195
+ if legacy:
196
+ #num_heads = 1
197
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
198
+ if exists(disable_self_attentions):
199
+ disabled_sa = disable_self_attentions[level]
200
+ else:
201
+ disabled_sa = False
202
+
203
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
204
+ layers.append(
205
+ SpatialTransformer(
206
+ ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim,
207
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
208
+ use_checkpoint=use_checkpoint, operations=operations
209
+ )
210
+ )
211
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
212
+ self.zero_convs.append(self.make_zero_conv(ch, operations=operations))
213
+ self._feature_size += ch
214
+ input_block_chans.append(ch)
215
+ if level != len(channel_mult) - 1:
216
+ out_ch = ch
217
+ self.input_blocks.append(
218
+ TimestepEmbedSequential(
219
+ ResBlock(
220
+ ch,
221
+ time_embed_dim,
222
+ dropout,
223
+ out_channels=out_ch,
224
+ dims=dims,
225
+ use_checkpoint=use_checkpoint,
226
+ use_scale_shift_norm=use_scale_shift_norm,
227
+ down=True,
228
+ operations=operations
229
+ )
230
+ if resblock_updown
231
+ else Downsample(
232
+ ch, conv_resample, dims=dims, out_channels=out_ch, operations=operations
233
+ )
234
+ )
235
+ )
236
+ ch = out_ch
237
+ input_block_chans.append(ch)
238
+ self.zero_convs.append(self.make_zero_conv(ch, operations=operations))
239
+ ds *= 2
240
+ self._feature_size += ch
241
+
242
+ if num_head_channels == -1:
243
+ dim_head = ch // num_heads
244
+ else:
245
+ num_heads = ch // num_head_channels
246
+ dim_head = num_head_channels
247
+ if legacy:
248
+ #num_heads = 1
249
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
250
+ self.middle_block = TimestepEmbedSequential(
251
+ ResBlock(
252
+ ch,
253
+ time_embed_dim,
254
+ dropout,
255
+ dims=dims,
256
+ use_checkpoint=use_checkpoint,
257
+ use_scale_shift_norm=use_scale_shift_norm,
258
+ operations=operations
259
+ ),
260
+ SpatialTransformer( # always uses a self-attn
261
+ ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
262
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
263
+ use_checkpoint=use_checkpoint, operations=operations
264
+ ),
265
+ ResBlock(
266
+ ch,
267
+ time_embed_dim,
268
+ dropout,
269
+ dims=dims,
270
+ use_checkpoint=use_checkpoint,
271
+ use_scale_shift_norm=use_scale_shift_norm,
272
+ operations=operations
273
+ ),
274
+ )
275
+ self.middle_block_out = self.make_zero_conv(ch, operations=operations)
276
+ self._feature_size += ch
277
+
278
+ def make_zero_conv(self, channels, operations=None):
279
+ return TimestepEmbedSequential(zero_module(operations.conv_nd(self.dims, channels, channels, 1, padding=0)))
280
+
281
+ def forward(self, x, hint, timesteps, context, y=None, **kwargs):
282
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(self.dtype)
283
+ emb = self.time_embed(t_emb)
284
+
285
+ guided_hint = self.input_hint_block(hint, emb, context)
286
+
287
+ outs = []
288
+
289
+ hs = []
290
+ if self.num_classes is not None:
291
+ assert y.shape[0] == x.shape[0]
292
+ emb = emb + self.label_emb(y)
293
+
294
+ h = x.type(self.dtype)
295
+ for module, zero_conv in zip(self.input_blocks, self.zero_convs):
296
+ if guided_hint is not None:
297
+ h = module(h, emb, context)
298
+ h += guided_hint
299
+ guided_hint = None
300
+ else:
301
+ h = module(h, emb, context)
302
+ outs.append(zero_conv(h, emb, context))
303
+
304
+ h = self.middle_block(h, emb, context)
305
+ outs.append(self.middle_block_out(h, emb, context))
306
+
307
+ return outs
308
+
comfy/cli_args.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import enum
3
+ import comfy.options
4
+
5
+ class EnumAction(argparse.Action):
6
+ """
7
+ Argparse action for handling Enums
8
+ """
9
+ def __init__(self, **kwargs):
10
+ # Pop off the type value
11
+ enum_type = kwargs.pop("type", None)
12
+
13
+ # Ensure an Enum subclass is provided
14
+ if enum_type is None:
15
+ raise ValueError("type must be assigned an Enum when using EnumAction")
16
+ if not issubclass(enum_type, enum.Enum):
17
+ raise TypeError("type must be an Enum when using EnumAction")
18
+
19
+ # Generate choices from the Enum
20
+ choices = tuple(e.value for e in enum_type)
21
+ kwargs.setdefault("choices", choices)
22
+ kwargs.setdefault("metavar", f"[{','.join(list(choices))}]")
23
+
24
+ super(EnumAction, self).__init__(**kwargs)
25
+
26
+ self._enum = enum_type
27
+
28
+ def __call__(self, parser, namespace, values, option_string=None):
29
+ # Convert value back into an Enum
30
+ value = self._enum(values)
31
+ setattr(namespace, self.dest, value)
32
+
33
+
34
+ parser = argparse.ArgumentParser()
35
+
36
+ parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0", help="Specify the IP address to listen on (default: 127.0.0.1). If --listen is provided without an argument, it defaults to 0.0.0.0. (listens on all)")
37
+ parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
38
+ parser.add_argument("--enable-cors-header", type=str, default=None, metavar="ORIGIN", nargs="?", const="*", help="Enable CORS (Cross-Origin Resource Sharing) with optional origin or allow all with default '*'.")
39
+ parser.add_argument("--extra-model-paths-config", type=str, default=None, metavar="PATH", nargs='+', action='append', help="Load one or more extra_model_paths.yaml files.")
40
+ parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
41
+ parser.add_argument("--temp-directory", type=str, default=None, help="Set the ComfyUI temp directory (default is in the ComfyUI directory).")
42
+ parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
43
+ parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
44
+ parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
45
+ cm_group = parser.add_mutually_exclusive_group()
46
+ cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
47
+ cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.")
48
+
49
+ parser.add_argument("--dont-upcast-attention", action="store_true", help="Disable upcasting of attention. Can boost speed but increase the chances of black images.")
50
+
51
+ fp_group = parser.add_mutually_exclusive_group()
52
+ fp_group.add_argument("--force-fp32", action="store_true", help="Force fp32 (If this makes your GPU work better please report it).")
53
+ fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.")
54
+
55
+ fpvae_group = parser.add_mutually_exclusive_group()
56
+ fpvae_group.add_argument("--fp16-vae", action="store_true", help="Run the VAE in fp16, might cause black images.")
57
+ fpvae_group.add_argument("--fp32-vae", action="store_true", help="Run the VAE in full precision fp32.")
58
+ fpvae_group.add_argument("--bf16-vae", action="store_true", help="Run the VAE in bf16.")
59
+
60
+ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1, help="Use torch-directml.")
61
+
62
+ parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize when loading models with Intel GPUs.")
63
+
64
+ class LatentPreviewMethod(enum.Enum):
65
+ NoPreviews = "none"
66
+ Auto = "auto"
67
+ Latent2RGB = "latent2rgb"
68
+ TAESD = "taesd"
69
+
70
+ parser.add_argument("--preview-method", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, help="Default preview method for sampler nodes.", action=EnumAction)
71
+
72
+ attn_group = parser.add_mutually_exclusive_group()
73
+ attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
74
+ attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.")
75
+ attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
76
+
77
+ parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
78
+
79
+ vram_group = parser.add_mutually_exclusive_group()
80
+ vram_group.add_argument("--gpu-only", action="store_true", help="Store and run everything (text encoders/CLIP models, etc... on the GPU).")
81
+ vram_group.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
82
+ vram_group.add_argument("--normalvram", action="store_true", help="Used to force normal vram use if lowvram gets automatically enabled.")
83
+ vram_group.add_argument("--lowvram", action="store_true", help="Split the unet in parts to use less vram.")
84
+ vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
85
+ vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
86
+
87
+
88
+ parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.")
89
+
90
+
91
+ parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
92
+ parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
93
+ parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).")
94
+
95
+ parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.")
96
+
97
+ if comfy.options.args_parsing:
98
+ args = parser.parse_args()
99
+ else:
100
+ args = parser.parse_args([])
101
+
102
+ if args.windows_standalone_build:
103
+ args.auto_launch = True
104
+
105
+ if args.disable_auto_launch:
106
+ args.auto_launch = False
comfy/clip_config_bigg.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "CLIPTextModel"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "bos_token_id": 0,
7
+ "dropout": 0.0,
8
+ "eos_token_id": 2,
9
+ "hidden_act": "gelu",
10
+ "hidden_size": 1280,
11
+ "initializer_factor": 1.0,
12
+ "initializer_range": 0.02,
13
+ "intermediate_size": 5120,
14
+ "layer_norm_eps": 1e-05,
15
+ "max_position_embeddings": 77,
16
+ "model_type": "clip_text_model",
17
+ "num_attention_heads": 20,
18
+ "num_hidden_layers": 32,
19
+ "pad_token_id": 1,
20
+ "projection_dim": 1280,
21
+ "torch_dtype": "float32",
22
+ "vocab_size": 49408
23
+ }
comfy/clip_vision.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CLIPVisionModelWithProjection, CLIPVisionConfig, CLIPImageProcessor, modeling_utils
2
+ from .utils import load_torch_file, transformers_convert
3
+ import os
4
+ import torch
5
+ import contextlib
6
+
7
+ import comfy.ops
8
+ import comfy.model_patcher
9
+ import comfy.model_management
10
+
11
+ class ClipVisionModel():
12
+ def __init__(self, json_config):
13
+ config = CLIPVisionConfig.from_json_file(json_config)
14
+ self.load_device = comfy.model_management.text_encoder_device()
15
+ offload_device = comfy.model_management.text_encoder_offload_device()
16
+ self.dtype = torch.float32
17
+ if comfy.model_management.should_use_fp16(self.load_device, prioritize_performance=False):
18
+ self.dtype = torch.float16
19
+
20
+ with comfy.ops.use_comfy_ops(offload_device, self.dtype):
21
+ with modeling_utils.no_init_weights():
22
+ self.model = CLIPVisionModelWithProjection(config)
23
+ self.model.to(self.dtype)
24
+
25
+ self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
26
+ self.processor = CLIPImageProcessor(crop_size=224,
27
+ do_center_crop=True,
28
+ do_convert_rgb=True,
29
+ do_normalize=True,
30
+ do_resize=True,
31
+ image_mean=[ 0.48145466,0.4578275,0.40821073],
32
+ image_std=[0.26862954,0.26130258,0.27577711],
33
+ resample=3, #bicubic
34
+ size=224)
35
+
36
+ def load_sd(self, sd):
37
+ return self.model.load_state_dict(sd, strict=False)
38
+
39
+ def encode_image(self, image):
40
+ img = torch.clip((255. * image), 0, 255).round().int()
41
+ img = list(map(lambda a: a, img))
42
+ inputs = self.processor(images=img, return_tensors="pt")
43
+ comfy.model_management.load_model_gpu(self.patcher)
44
+ pixel_values = inputs['pixel_values'].to(self.load_device)
45
+
46
+ if self.dtype != torch.float32:
47
+ precision_scope = torch.autocast
48
+ else:
49
+ precision_scope = lambda a, b: contextlib.nullcontext(a)
50
+
51
+ with precision_scope(comfy.model_management.get_autocast_device(self.load_device), torch.float32):
52
+ outputs = self.model(pixel_values=pixel_values, output_hidden_states=True)
53
+
54
+ for k in outputs:
55
+ t = outputs[k]
56
+ if t is not None:
57
+ if k == 'hidden_states':
58
+ outputs["penultimate_hidden_states"] = t[-2].cpu()
59
+ outputs["hidden_states"] = None
60
+ else:
61
+ outputs[k] = t.cpu()
62
+
63
+ return outputs
64
+
65
+ def convert_to_transformers(sd, prefix):
66
+ sd_k = sd.keys()
67
+ if "{}transformer.resblocks.0.attn.in_proj_weight".format(prefix) in sd_k:
68
+ keys_to_replace = {
69
+ "{}class_embedding".format(prefix): "vision_model.embeddings.class_embedding",
70
+ "{}conv1.weight".format(prefix): "vision_model.embeddings.patch_embedding.weight",
71
+ "{}positional_embedding".format(prefix): "vision_model.embeddings.position_embedding.weight",
72
+ "{}ln_post.bias".format(prefix): "vision_model.post_layernorm.bias",
73
+ "{}ln_post.weight".format(prefix): "vision_model.post_layernorm.weight",
74
+ "{}ln_pre.bias".format(prefix): "vision_model.pre_layrnorm.bias",
75
+ "{}ln_pre.weight".format(prefix): "vision_model.pre_layrnorm.weight",
76
+ }
77
+
78
+ for x in keys_to_replace:
79
+ if x in sd_k:
80
+ sd[keys_to_replace[x]] = sd.pop(x)
81
+
82
+ if "{}proj".format(prefix) in sd_k:
83
+ sd['visual_projection.weight'] = sd.pop("{}proj".format(prefix)).transpose(0, 1)
84
+
85
+ sd = transformers_convert(sd, prefix, "vision_model.", 48)
86
+ return sd
87
+
88
+ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
89
+ if convert_keys:
90
+ sd = convert_to_transformers(sd, prefix)
91
+ if "vision_model.encoder.layers.47.layer_norm1.weight" in sd:
92
+ json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_g.json")
93
+ elif "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
94
+ json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_h.json")
95
+ else:
96
+ json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
97
+ clip = ClipVisionModel(json_config)
98
+ m, u = clip.load_sd(sd)
99
+ if len(m) > 0:
100
+ print("missing clip vision:", m)
101
+ u = set(u)
102
+ keys = list(sd.keys())
103
+ for k in keys:
104
+ if k not in u:
105
+ t = sd.pop(k)
106
+ del t
107
+ return clip
108
+
109
+ def load(ckpt_path):
110
+ sd = load_torch_file(ckpt_path)
111
+ if "visual.transformer.resblocks.0.attn.in_proj_weight" in sd:
112
+ return load_clipvision_from_sd(sd, prefix="visual.", convert_keys=True)
113
+ else:
114
+ return load_clipvision_from_sd(sd)
comfy/clip_vision_config_g.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "attention_dropout": 0.0,
3
+ "dropout": 0.0,
4
+ "hidden_act": "gelu",
5
+ "hidden_size": 1664,
6
+ "image_size": 224,
7
+ "initializer_factor": 1.0,
8
+ "initializer_range": 0.02,
9
+ "intermediate_size": 8192,
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": 48,
15
+ "patch_size": 14,
16
+ "projection_dim": 1280,
17
+ "torch_dtype": "float32"
18
+ }
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/controlnet.py ADDED
@@ -0,0 +1,488 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import math
3
+ import os
4
+ import comfy.utils
5
+ import comfy.model_management
6
+ import comfy.model_detection
7
+ import comfy.model_patcher
8
+
9
+ import comfy.cldm.cldm
10
+ import comfy.t2i_adapter.adapter
11
+
12
+
13
+ def broadcast_image_to(tensor, target_batch_size, batched_number):
14
+ current_batch_size = tensor.shape[0]
15
+ #print(current_batch_size, target_batch_size)
16
+ if current_batch_size == 1:
17
+ return tensor
18
+
19
+ per_batch = target_batch_size // batched_number
20
+ tensor = tensor[:per_batch]
21
+
22
+ if per_batch > tensor.shape[0]:
23
+ tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0)
24
+
25
+ current_batch_size = tensor.shape[0]
26
+ if current_batch_size == target_batch_size:
27
+ return tensor
28
+ else:
29
+ return torch.cat([tensor] * batched_number, dim=0)
30
+
31
+ class ControlBase:
32
+ def __init__(self, device=None):
33
+ self.cond_hint_original = None
34
+ self.cond_hint = None
35
+ self.strength = 1.0
36
+ self.timestep_percent_range = (1.0, 0.0)
37
+ self.timestep_range = None
38
+
39
+ if device is None:
40
+ device = comfy.model_management.get_torch_device()
41
+ self.device = device
42
+ self.previous_controlnet = None
43
+ self.global_average_pooling = False
44
+
45
+ def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(1.0, 0.0)):
46
+ self.cond_hint_original = cond_hint
47
+ self.strength = strength
48
+ self.timestep_percent_range = timestep_percent_range
49
+ return self
50
+
51
+ def pre_run(self, model, percent_to_timestep_function):
52
+ self.timestep_range = (percent_to_timestep_function(self.timestep_percent_range[0]), percent_to_timestep_function(self.timestep_percent_range[1]))
53
+ if self.previous_controlnet is not None:
54
+ self.previous_controlnet.pre_run(model, percent_to_timestep_function)
55
+
56
+ def set_previous_controlnet(self, controlnet):
57
+ self.previous_controlnet = controlnet
58
+ return self
59
+
60
+ def cleanup(self):
61
+ if self.previous_controlnet is not None:
62
+ self.previous_controlnet.cleanup()
63
+ if self.cond_hint is not None:
64
+ del self.cond_hint
65
+ self.cond_hint = None
66
+ self.timestep_range = None
67
+
68
+ def get_models(self):
69
+ out = []
70
+ if self.previous_controlnet is not None:
71
+ out += self.previous_controlnet.get_models()
72
+ return out
73
+
74
+ def copy_to(self, c):
75
+ c.cond_hint_original = self.cond_hint_original
76
+ c.strength = self.strength
77
+ c.timestep_percent_range = self.timestep_percent_range
78
+
79
+ def inference_memory_requirements(self, dtype):
80
+ if self.previous_controlnet is not None:
81
+ return self.previous_controlnet.inference_memory_requirements(dtype)
82
+ return 0
83
+
84
+ def control_merge(self, control_input, control_output, control_prev, output_dtype):
85
+ out = {'input':[], 'middle':[], 'output': []}
86
+
87
+ if control_input is not None:
88
+ for i in range(len(control_input)):
89
+ key = 'input'
90
+ x = control_input[i]
91
+ if x is not None:
92
+ x *= self.strength
93
+ if x.dtype != output_dtype:
94
+ x = x.to(output_dtype)
95
+ out[key].insert(0, x)
96
+
97
+ if control_output is not None:
98
+ for i in range(len(control_output)):
99
+ if i == (len(control_output) - 1):
100
+ key = 'middle'
101
+ index = 0
102
+ else:
103
+ key = 'output'
104
+ index = i
105
+ x = control_output[i]
106
+ if x is not None:
107
+ if self.global_average_pooling:
108
+ x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
109
+
110
+ x *= self.strength
111
+ if x.dtype != output_dtype:
112
+ x = x.to(output_dtype)
113
+
114
+ out[key].append(x)
115
+ if control_prev is not None:
116
+ for x in ['input', 'middle', 'output']:
117
+ o = out[x]
118
+ for i in range(len(control_prev[x])):
119
+ prev_val = control_prev[x][i]
120
+ if i >= len(o):
121
+ o.append(prev_val)
122
+ elif prev_val is not None:
123
+ if o[i] is None:
124
+ o[i] = prev_val
125
+ else:
126
+ o[i] += prev_val
127
+ return out
128
+
129
+ class ControlNet(ControlBase):
130
+ def __init__(self, control_model, global_average_pooling=False, device=None):
131
+ super().__init__(device)
132
+ self.control_model = control_model
133
+ self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
134
+ self.global_average_pooling = global_average_pooling
135
+
136
+ def get_control(self, x_noisy, t, cond, batched_number):
137
+ control_prev = None
138
+ if self.previous_controlnet is not None:
139
+ control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
140
+
141
+ if self.timestep_range is not None:
142
+ if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
143
+ if control_prev is not None:
144
+ return control_prev
145
+ else:
146
+ return None
147
+
148
+ output_dtype = x_noisy.dtype
149
+ if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
150
+ if self.cond_hint is not None:
151
+ del self.cond_hint
152
+ self.cond_hint = None
153
+ self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(self.control_model.dtype).to(self.device)
154
+ if x_noisy.shape[0] != self.cond_hint.shape[0]:
155
+ self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
156
+
157
+
158
+ context = cond['c_crossattn']
159
+ y = cond.get('c_adm', None)
160
+ if y is not None:
161
+ y = y.to(self.control_model.dtype)
162
+ control = self.control_model(x=x_noisy.to(self.control_model.dtype), hint=self.cond_hint, timesteps=t, context=context.to(self.control_model.dtype), y=y)
163
+ return self.control_merge(None, control, control_prev, output_dtype)
164
+
165
+ def copy(self):
166
+ c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling)
167
+ self.copy_to(c)
168
+ return c
169
+
170
+ def get_models(self):
171
+ out = super().get_models()
172
+ out.append(self.control_model_wrapped)
173
+ return out
174
+
175
+ class ControlLoraOps:
176
+ class Linear(torch.nn.Module):
177
+ def __init__(self, in_features: int, out_features: int, bias: bool = True,
178
+ device=None, dtype=None) -> None:
179
+ factory_kwargs = {'device': device, 'dtype': dtype}
180
+ super().__init__()
181
+ self.in_features = in_features
182
+ self.out_features = out_features
183
+ self.weight = None
184
+ self.up = None
185
+ self.down = None
186
+ self.bias = None
187
+
188
+ def forward(self, input):
189
+ if self.up is not None:
190
+ return torch.nn.functional.linear(input, self.weight.to(input.device) + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), self.bias)
191
+ else:
192
+ return torch.nn.functional.linear(input, self.weight.to(input.device), self.bias)
193
+
194
+ class Conv2d(torch.nn.Module):
195
+ def __init__(
196
+ self,
197
+ in_channels,
198
+ out_channels,
199
+ kernel_size,
200
+ stride=1,
201
+ padding=0,
202
+ dilation=1,
203
+ groups=1,
204
+ bias=True,
205
+ padding_mode='zeros',
206
+ device=None,
207
+ dtype=None
208
+ ):
209
+ super().__init__()
210
+ self.in_channels = in_channels
211
+ self.out_channels = out_channels
212
+ self.kernel_size = kernel_size
213
+ self.stride = stride
214
+ self.padding = padding
215
+ self.dilation = dilation
216
+ self.transposed = False
217
+ self.output_padding = 0
218
+ self.groups = groups
219
+ self.padding_mode = padding_mode
220
+
221
+ self.weight = None
222
+ self.bias = None
223
+ self.up = None
224
+ self.down = None
225
+
226
+
227
+ def forward(self, input):
228
+ if self.up is not None:
229
+ return torch.nn.functional.conv2d(input, self.weight.to(input.device) + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), self.bias, self.stride, self.padding, self.dilation, self.groups)
230
+ else:
231
+ return torch.nn.functional.conv2d(input, self.weight.to(input.device), self.bias, self.stride, self.padding, self.dilation, self.groups)
232
+
233
+ def conv_nd(self, dims, *args, **kwargs):
234
+ if dims == 2:
235
+ return self.Conv2d(*args, **kwargs)
236
+ else:
237
+ raise ValueError(f"unsupported dimensions: {dims}")
238
+
239
+
240
+ class ControlLora(ControlNet):
241
+ def __init__(self, control_weights, global_average_pooling=False, device=None):
242
+ ControlBase.__init__(self, device)
243
+ self.control_weights = control_weights
244
+ self.global_average_pooling = global_average_pooling
245
+
246
+ def pre_run(self, model, percent_to_timestep_function):
247
+ super().pre_run(model, percent_to_timestep_function)
248
+ controlnet_config = model.model_config.unet_config.copy()
249
+ controlnet_config.pop("out_channels")
250
+ controlnet_config["hint_channels"] = self.control_weights["input_hint_block.0.weight"].shape[1]
251
+ controlnet_config["operations"] = ControlLoraOps()
252
+ self.control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)
253
+ dtype = model.get_dtype()
254
+ self.control_model.to(dtype)
255
+ self.control_model.to(comfy.model_management.get_torch_device())
256
+ diffusion_model = model.diffusion_model
257
+ sd = diffusion_model.state_dict()
258
+ cm = self.control_model.state_dict()
259
+
260
+ for k in sd:
261
+ weight = comfy.model_management.resolve_lowvram_weight(sd[k], diffusion_model, k)
262
+ try:
263
+ comfy.utils.set_attr(self.control_model, k, weight)
264
+ except:
265
+ pass
266
+
267
+ for k in self.control_weights:
268
+ if k not in {"lora_controlnet"}:
269
+ comfy.utils.set_attr(self.control_model, k, self.control_weights[k].to(dtype).to(comfy.model_management.get_torch_device()))
270
+
271
+ def copy(self):
272
+ c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)
273
+ self.copy_to(c)
274
+ return c
275
+
276
+ def cleanup(self):
277
+ del self.control_model
278
+ self.control_model = None
279
+ super().cleanup()
280
+
281
+ def get_models(self):
282
+ out = ControlBase.get_models(self)
283
+ return out
284
+
285
+ def inference_memory_requirements(self, dtype):
286
+ return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
287
+
288
+ def load_controlnet(ckpt_path, model=None):
289
+ controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
290
+ if "lora_controlnet" in controlnet_data:
291
+ return ControlLora(controlnet_data)
292
+
293
+ controlnet_config = None
294
+ if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
295
+ use_fp16 = comfy.model_management.should_use_fp16()
296
+ controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data, use_fp16)
297
+ diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config)
298
+ diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
299
+ diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
300
+
301
+ count = 0
302
+ loop = True
303
+ while loop:
304
+ suffix = [".weight", ".bias"]
305
+ for s in suffix:
306
+ k_in = "controlnet_down_blocks.{}{}".format(count, s)
307
+ k_out = "zero_convs.{}.0{}".format(count, s)
308
+ if k_in not in controlnet_data:
309
+ loop = False
310
+ break
311
+ diffusers_keys[k_in] = k_out
312
+ count += 1
313
+
314
+ count = 0
315
+ loop = True
316
+ while loop:
317
+ suffix = [".weight", ".bias"]
318
+ for s in suffix:
319
+ if count == 0:
320
+ k_in = "controlnet_cond_embedding.conv_in{}".format(s)
321
+ else:
322
+ k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
323
+ k_out = "input_hint_block.{}{}".format(count * 2, s)
324
+ if k_in not in controlnet_data:
325
+ k_in = "controlnet_cond_embedding.conv_out{}".format(s)
326
+ loop = False
327
+ diffusers_keys[k_in] = k_out
328
+ count += 1
329
+
330
+ new_sd = {}
331
+ for k in diffusers_keys:
332
+ if k in controlnet_data:
333
+ new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
334
+
335
+ leftover_keys = controlnet_data.keys()
336
+ if len(leftover_keys) > 0:
337
+ print("leftover keys:", leftover_keys)
338
+ controlnet_data = new_sd
339
+
340
+ pth_key = 'control_model.zero_convs.0.0.weight'
341
+ pth = False
342
+ key = 'zero_convs.0.0.weight'
343
+ if pth_key in controlnet_data:
344
+ pth = True
345
+ key = pth_key
346
+ prefix = "control_model."
347
+ elif key in controlnet_data:
348
+ prefix = ""
349
+ else:
350
+ net = load_t2i_adapter(controlnet_data)
351
+ if net is None:
352
+ print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path)
353
+ return net
354
+
355
+ if controlnet_config is None:
356
+ use_fp16 = comfy.model_management.should_use_fp16()
357
+ controlnet_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, use_fp16, True).unet_config
358
+ controlnet_config.pop("out_channels")
359
+ controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
360
+ control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)
361
+
362
+ if pth:
363
+ if 'difference' in controlnet_data:
364
+ if model is not None:
365
+ comfy.model_management.load_models_gpu([model])
366
+ model_sd = model.model_state_dict()
367
+ for x in controlnet_data:
368
+ c_m = "control_model."
369
+ if x.startswith(c_m):
370
+ sd_key = "diffusion_model.{}".format(x[len(c_m):])
371
+ if sd_key in model_sd:
372
+ cd = controlnet_data[x]
373
+ cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
374
+ else:
375
+ print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
376
+
377
+ class WeightsLoader(torch.nn.Module):
378
+ pass
379
+ w = WeightsLoader()
380
+ w.control_model = control_model
381
+ missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
382
+ else:
383
+ missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
384
+ print(missing, unexpected)
385
+
386
+ if use_fp16:
387
+ control_model = control_model.half()
388
+
389
+ global_average_pooling = False
390
+ filename = os.path.splitext(ckpt_path)[0]
391
+ if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
392
+ global_average_pooling = True
393
+
394
+ control = ControlNet(control_model, global_average_pooling=global_average_pooling)
395
+ return control
396
+
397
+ class T2IAdapter(ControlBase):
398
+ def __init__(self, t2i_model, channels_in, device=None):
399
+ super().__init__(device)
400
+ self.t2i_model = t2i_model
401
+ self.channels_in = channels_in
402
+ self.control_input = None
403
+
404
+ def scale_image_to(self, width, height):
405
+ unshuffle_amount = self.t2i_model.unshuffle_amount
406
+ width = math.ceil(width / unshuffle_amount) * unshuffle_amount
407
+ height = math.ceil(height / unshuffle_amount) * unshuffle_amount
408
+ return width, height
409
+
410
+ def get_control(self, x_noisy, t, cond, batched_number):
411
+ control_prev = None
412
+ if self.previous_controlnet is not None:
413
+ control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
414
+
415
+ if self.timestep_range is not None:
416
+ if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
417
+ if control_prev is not None:
418
+ return control_prev
419
+ else:
420
+ return {}
421
+
422
+ if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
423
+ if self.cond_hint is not None:
424
+ del self.cond_hint
425
+ self.control_input = None
426
+ self.cond_hint = None
427
+ width, height = self.scale_image_to(x_noisy.shape[3] * 8, x_noisy.shape[2] * 8)
428
+ self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, width, height, 'nearest-exact', "center").float().to(self.device)
429
+ if self.channels_in == 1 and self.cond_hint.shape[1] > 1:
430
+ self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)
431
+ if x_noisy.shape[0] != self.cond_hint.shape[0]:
432
+ self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
433
+ if self.control_input is None:
434
+ self.t2i_model.to(x_noisy.dtype)
435
+ self.t2i_model.to(self.device)
436
+ self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
437
+ self.t2i_model.cpu()
438
+
439
+ control_input = list(map(lambda a: None if a is None else a.clone(), self.control_input))
440
+ mid = None
441
+ if self.t2i_model.xl == True:
442
+ mid = control_input[-1:]
443
+ control_input = control_input[:-1]
444
+ return self.control_merge(control_input, mid, control_prev, x_noisy.dtype)
445
+
446
+ def copy(self):
447
+ c = T2IAdapter(self.t2i_model, self.channels_in)
448
+ self.copy_to(c)
449
+ return c
450
+
451
+ def load_t2i_adapter(t2i_data):
452
+ if 'adapter' in t2i_data:
453
+ t2i_data = t2i_data['adapter']
454
+ if 'adapter.body.0.resnets.0.block1.weight' in t2i_data: #diffusers format
455
+ prefix_replace = {}
456
+ for i in range(4):
457
+ for j in range(2):
458
+ prefix_replace["adapter.body.{}.resnets.{}.".format(i, j)] = "body.{}.".format(i * 2 + j)
459
+ prefix_replace["adapter.body.{}.".format(i, j)] = "body.{}.".format(i * 2)
460
+ prefix_replace["adapter."] = ""
461
+ t2i_data = comfy.utils.state_dict_prefix_replace(t2i_data, prefix_replace)
462
+ keys = t2i_data.keys()
463
+
464
+ if "body.0.in_conv.weight" in keys:
465
+ cin = t2i_data['body.0.in_conv.weight'].shape[1]
466
+ model_ad = comfy.t2i_adapter.adapter.Adapter_light(cin=cin, channels=[320, 640, 1280, 1280], nums_rb=4)
467
+ elif 'conv_in.weight' in keys:
468
+ cin = t2i_data['conv_in.weight'].shape[1]
469
+ channel = t2i_data['conv_in.weight'].shape[0]
470
+ ksize = t2i_data['body.0.block2.weight'].shape[2]
471
+ use_conv = False
472
+ down_opts = list(filter(lambda a: a.endswith("down_opt.op.weight"), keys))
473
+ if len(down_opts) > 0:
474
+ use_conv = True
475
+ xl = False
476
+ if cin == 256 or cin == 768:
477
+ xl = True
478
+ model_ad = comfy.t2i_adapter.adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)
479
+ else:
480
+ return None
481
+ missing, unexpected = model_ad.load_state_dict(t2i_data)
482
+ if len(missing) > 0:
483
+ print("t2i missing", missing)
484
+
485
+ if len(unexpected) > 0:
486
+ print("t2i unexpected", unexpected)
487
+
488
+ return T2IAdapter(model_ad, model_ad.input_channels)
comfy/diffusers_convert.py ADDED
@@ -0,0 +1,261 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import torch
3
+
4
+ # conversion code from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
5
+
6
+ # =================#
7
+ # UNet Conversion #
8
+ # =================#
9
+
10
+ unet_conversion_map = [
11
+ # (stable-diffusion, HF Diffusers)
12
+ ("time_embed.0.weight", "time_embedding.linear_1.weight"),
13
+ ("time_embed.0.bias", "time_embedding.linear_1.bias"),
14
+ ("time_embed.2.weight", "time_embedding.linear_2.weight"),
15
+ ("time_embed.2.bias", "time_embedding.linear_2.bias"),
16
+ ("input_blocks.0.0.weight", "conv_in.weight"),
17
+ ("input_blocks.0.0.bias", "conv_in.bias"),
18
+ ("out.0.weight", "conv_norm_out.weight"),
19
+ ("out.0.bias", "conv_norm_out.bias"),
20
+ ("out.2.weight", "conv_out.weight"),
21
+ ("out.2.bias", "conv_out.bias"),
22
+ ]
23
+
24
+ unet_conversion_map_resnet = [
25
+ # (stable-diffusion, HF Diffusers)
26
+ ("in_layers.0", "norm1"),
27
+ ("in_layers.2", "conv1"),
28
+ ("out_layers.0", "norm2"),
29
+ ("out_layers.3", "conv2"),
30
+ ("emb_layers.1", "time_emb_proj"),
31
+ ("skip_connection", "conv_shortcut"),
32
+ ]
33
+
34
+ unet_conversion_map_layer = []
35
+ # hardcoded number of downblocks and resnets/attentions...
36
+ # would need smarter logic for other networks.
37
+ for i in range(4):
38
+ # loop over downblocks/upblocks
39
+
40
+ for j in range(2):
41
+ # loop over resnets/attentions for downblocks
42
+ hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
43
+ sd_down_res_prefix = f"input_blocks.{3 * i + j + 1}.0."
44
+ unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
45
+
46
+ if i < 3:
47
+ # no attention layers in down_blocks.3
48
+ hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
49
+ sd_down_atn_prefix = f"input_blocks.{3 * i + j + 1}.1."
50
+ unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
51
+
52
+ for j in range(3):
53
+ # loop over resnets/attentions for upblocks
54
+ hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
55
+ sd_up_res_prefix = f"output_blocks.{3 * i + j}.0."
56
+ unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
57
+
58
+ if i > 0:
59
+ # no attention layers in up_blocks.0
60
+ hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
61
+ sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1."
62
+ unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
63
+
64
+ if i < 3:
65
+ # no downsample in down_blocks.3
66
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
67
+ sd_downsample_prefix = f"input_blocks.{3 * (i + 1)}.0.op."
68
+ unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
69
+
70
+ # no upsample in up_blocks.3
71
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
72
+ sd_upsample_prefix = f"output_blocks.{3 * i + 2}.{1 if i == 0 else 2}."
73
+ unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
74
+
75
+ hf_mid_atn_prefix = "mid_block.attentions.0."
76
+ sd_mid_atn_prefix = "middle_block.1."
77
+ unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
78
+
79
+ for j in range(2):
80
+ hf_mid_res_prefix = f"mid_block.resnets.{j}."
81
+ sd_mid_res_prefix = f"middle_block.{2 * j}."
82
+ unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
83
+
84
+
85
+ def convert_unet_state_dict(unet_state_dict):
86
+ # buyer beware: this is a *brittle* function,
87
+ # and correct output requires that all of these pieces interact in
88
+ # the exact order in which I have arranged them.
89
+ mapping = {k: k for k in unet_state_dict.keys()}
90
+ for sd_name, hf_name in unet_conversion_map:
91
+ mapping[hf_name] = sd_name
92
+ for k, v in mapping.items():
93
+ if "resnets" in k:
94
+ for sd_part, hf_part in unet_conversion_map_resnet:
95
+ v = v.replace(hf_part, sd_part)
96
+ mapping[k] = v
97
+ for k, v in mapping.items():
98
+ for sd_part, hf_part in unet_conversion_map_layer:
99
+ v = v.replace(hf_part, sd_part)
100
+ mapping[k] = v
101
+ new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
102
+ return new_state_dict
103
+
104
+
105
+ # ================#
106
+ # VAE Conversion #
107
+ # ================#
108
+
109
+ vae_conversion_map = [
110
+ # (stable-diffusion, HF Diffusers)
111
+ ("nin_shortcut", "conv_shortcut"),
112
+ ("norm_out", "conv_norm_out"),
113
+ ("mid.attn_1.", "mid_block.attentions.0."),
114
+ ]
115
+
116
+ for i in range(4):
117
+ # down_blocks have two resnets
118
+ for j in range(2):
119
+ hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
120
+ sd_down_prefix = f"encoder.down.{i}.block.{j}."
121
+ vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
122
+
123
+ if i < 3:
124
+ hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
125
+ sd_downsample_prefix = f"down.{i}.downsample."
126
+ vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
127
+
128
+ hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
129
+ sd_upsample_prefix = f"up.{3 - i}.upsample."
130
+ vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
131
+
132
+ # up_blocks have three resnets
133
+ # also, up blocks in hf are numbered in reverse from sd
134
+ for j in range(3):
135
+ hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
136
+ sd_up_prefix = f"decoder.up.{3 - i}.block.{j}."
137
+ vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
138
+
139
+ # this part accounts for mid blocks in both the encoder and the decoder
140
+ for i in range(2):
141
+ hf_mid_res_prefix = f"mid_block.resnets.{i}."
142
+ sd_mid_res_prefix = f"mid.block_{i + 1}."
143
+ vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
144
+
145
+ vae_conversion_map_attn = [
146
+ # (stable-diffusion, HF Diffusers)
147
+ ("norm.", "group_norm."),
148
+ ("q.", "query."),
149
+ ("k.", "key."),
150
+ ("v.", "value."),
151
+ ("q.", "to_q."),
152
+ ("k.", "to_k."),
153
+ ("v.", "to_v."),
154
+ ("proj_out.", "to_out.0."),
155
+ ("proj_out.", "proj_attn."),
156
+ ]
157
+
158
+
159
+ def reshape_weight_for_sd(w):
160
+ # convert HF linear weights to SD conv2d weights
161
+ return w.reshape(*w.shape, 1, 1)
162
+
163
+
164
+ def convert_vae_state_dict(vae_state_dict):
165
+ mapping = {k: k for k in vae_state_dict.keys()}
166
+ for k, v in mapping.items():
167
+ for sd_part, hf_part in vae_conversion_map:
168
+ v = v.replace(hf_part, sd_part)
169
+ mapping[k] = v
170
+ for k, v in mapping.items():
171
+ if "attentions" in k:
172
+ for sd_part, hf_part in vae_conversion_map_attn:
173
+ v = v.replace(hf_part, sd_part)
174
+ mapping[k] = v
175
+ new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
176
+ weights_to_convert = ["q", "k", "v", "proj_out"]
177
+ for k, v in new_state_dict.items():
178
+ for weight_name in weights_to_convert:
179
+ if f"mid.attn_1.{weight_name}.weight" in k:
180
+ print(f"Reshaping {k} for SD format")
181
+ new_state_dict[k] = reshape_weight_for_sd(v)
182
+ return new_state_dict
183
+
184
+
185
+ # =========================#
186
+ # Text Encoder Conversion #
187
+ # =========================#
188
+
189
+
190
+ textenc_conversion_lst = [
191
+ # (stable-diffusion, HF Diffusers)
192
+ ("resblocks.", "text_model.encoder.layers."),
193
+ ("ln_1", "layer_norm1"),
194
+ ("ln_2", "layer_norm2"),
195
+ (".c_fc.", ".fc1."),
196
+ (".c_proj.", ".fc2."),
197
+ (".attn", ".self_attn"),
198
+ ("ln_final.", "transformer.text_model.final_layer_norm."),
199
+ ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
200
+ ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
201
+ ]
202
+ protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
203
+ textenc_pattern = re.compile("|".join(protected.keys()))
204
+
205
+ # Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
206
+ code2idx = {"q": 0, "k": 1, "v": 2}
207
+
208
+
209
+ def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""):
210
+ new_state_dict = {}
211
+ capture_qkv_weight = {}
212
+ capture_qkv_bias = {}
213
+ for k, v in text_enc_dict.items():
214
+ if not k.startswith(prefix):
215
+ continue
216
+ if (
217
+ k.endswith(".self_attn.q_proj.weight")
218
+ or k.endswith(".self_attn.k_proj.weight")
219
+ or k.endswith(".self_attn.v_proj.weight")
220
+ ):
221
+ k_pre = k[: -len(".q_proj.weight")]
222
+ k_code = k[-len("q_proj.weight")]
223
+ if k_pre not in capture_qkv_weight:
224
+ capture_qkv_weight[k_pre] = [None, None, None]
225
+ capture_qkv_weight[k_pre][code2idx[k_code]] = v
226
+ continue
227
+
228
+ if (
229
+ k.endswith(".self_attn.q_proj.bias")
230
+ or k.endswith(".self_attn.k_proj.bias")
231
+ or k.endswith(".self_attn.v_proj.bias")
232
+ ):
233
+ k_pre = k[: -len(".q_proj.bias")]
234
+ k_code = k[-len("q_proj.bias")]
235
+ if k_pre not in capture_qkv_bias:
236
+ capture_qkv_bias[k_pre] = [None, None, None]
237
+ capture_qkv_bias[k_pre][code2idx[k_code]] = v
238
+ continue
239
+
240
+ relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
241
+ new_state_dict[relabelled_key] = v
242
+
243
+ for k_pre, tensors in capture_qkv_weight.items():
244
+ if None in tensors:
245
+ raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
246
+ relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
247
+ new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
248
+
249
+ for k_pre, tensors in capture_qkv_bias.items():
250
+ if None in tensors:
251
+ raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
252
+ relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
253
+ new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
254
+
255
+ return new_state_dict
256
+
257
+
258
+ def convert_text_enc_state_dict(text_enc_dict):
259
+ return text_enc_dict
260
+
261
+
comfy/diffusers_load.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import os
3
+
4
+ import comfy.sd
5
+
6
+ def first_file(path, filenames):
7
+ for f in filenames:
8
+ p = os.path.join(path, f)
9
+ if os.path.exists(p):
10
+ return p
11
+ return None
12
+
13
+ def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None):
14
+ diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"]
15
+ unet_path = first_file(os.path.join(model_path, "unet"), diffusion_model_names)
16
+ vae_path = first_file(os.path.join(model_path, "vae"), diffusion_model_names)
17
+
18
+ text_encoder_model_names = ["model.fp16.safetensors", "model.safetensors", "pytorch_model.fp16.bin", "pytorch_model.bin"]
19
+ text_encoder1_path = first_file(os.path.join(model_path, "text_encoder"), text_encoder_model_names)
20
+ text_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)
21
+
22
+ text_encoder_paths = [text_encoder1_path]
23
+ if text_encoder2_path is not None:
24
+ text_encoder_paths.append(text_encoder2_path)
25
+
26
+ unet = comfy.sd.load_unet(unet_path)
27
+
28
+ clip = None
29
+ if output_clip:
30
+ clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory)
31
+
32
+ vae = None
33
+ if output_vae:
34
+ vae = comfy.sd.VAE(ckpt_path=vae_path)
35
+
36
+ return (unet, clip, vae)
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
+ noise_schedule,
184
+ model_type="noise",
185
+ model_kwargs={},
186
+ guidance_type="uncond",
187
+ condition=None,
188
+ unconditional_condition=None,
189
+ guidance_scale=1.,
190
+ classifier_fn=None,
191
+ classifier_kwargs={},
192
+ ):
193
+ """Create a wrapper function for the noise prediction model.
194
+
195
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
196
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
197
+
198
+ We support four types of the diffusion model by setting `model_type`:
199
+
200
+ 1. "noise": noise prediction model. (Trained by predicting noise).
201
+
202
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
203
+
204
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
205
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
206
+
207
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
208
+ arXiv preprint arXiv:2202.00512 (2022).
209
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
210
+ arXiv preprint arXiv:2210.02303 (2022).
211
+
212
+ 4. "score": marginal score function. (Trained by denoising score matching).
213
+ Note that the score function and the noise prediction model follows a simple relationship:
214
+ ```
215
+ noise(x_t, t) = -sigma_t * score(x_t, t)
216
+ ```
217
+
218
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
219
+ 1. "uncond": unconditional sampling by DPMs.
220
+ The input `model` has the following format:
221
+ ``
222
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
223
+ ``
224
+
225
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
226
+ The input `model` has the following format:
227
+ ``
228
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
229
+ ``
230
+
231
+ The input `classifier_fn` has the following format:
232
+ ``
233
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
234
+ ``
235
+
236
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
237
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
238
+
239
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
240
+ The input `model` has the following format:
241
+ ``
242
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
243
+ ``
244
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
245
+
246
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
247
+ arXiv preprint arXiv:2207.12598 (2022).
248
+
249
+
250
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
251
+ or continuous-time labels (i.e. epsilon to T).
252
+
253
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
254
+ ``
255
+ def model_fn(x, t_continuous) -> noise:
256
+ t_input = get_model_input_time(t_continuous)
257
+ return noise_pred(model, x, t_input, **model_kwargs)
258
+ ``
259
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
260
+
261
+ ===============================================================
262
+
263
+ Args:
264
+ model: A diffusion model with the corresponding format described above.
265
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
266
+ model_type: A `str`. The parameterization type of the diffusion model.
267
+ "noise" or "x_start" or "v" or "score".
268
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
269
+ guidance_type: A `str`. The type of the guidance for sampling.
270
+ "uncond" or "classifier" or "classifier-free".
271
+ condition: A pytorch tensor. The condition for the guided sampling.
272
+ Only used for "classifier" or "classifier-free" guidance type.
273
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
274
+ Only used for "classifier-free" guidance type.
275
+ guidance_scale: A `float`. The scale for the guided sampling.
276
+ classifier_fn: A classifier function. Only used for the classifier guidance.
277
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
278
+ Returns:
279
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
280
+ """
281
+
282
+ def get_model_input_time(t_continuous):
283
+ """
284
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
285
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
286
+ For continuous-time DPMs, we just use `t_continuous`.
287
+ """
288
+ if noise_schedule.schedule == 'discrete':
289
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
290
+ else:
291
+ return t_continuous
292
+
293
+ def noise_pred_fn(x, t_continuous, cond=None):
294
+ if t_continuous.reshape((-1,)).shape[0] == 1:
295
+ t_continuous = t_continuous.expand((x.shape[0]))
296
+ t_input = get_model_input_time(t_continuous)
297
+ output = model(x, t_input, **model_kwargs)
298
+ if model_type == "noise":
299
+ return output
300
+ elif model_type == "x_start":
301
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
302
+ dims = x.dim()
303
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
304
+ elif model_type == "v":
305
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
306
+ dims = x.dim()
307
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
308
+ elif model_type == "score":
309
+ sigma_t = noise_schedule.marginal_std(t_continuous)
310
+ dims = x.dim()
311
+ return -expand_dims(sigma_t, dims) * output
312
+
313
+ def cond_grad_fn(x, t_input):
314
+ """
315
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
316
+ """
317
+ with torch.enable_grad():
318
+ x_in = x.detach().requires_grad_(True)
319
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
320
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
321
+
322
+ def model_fn(x, t_continuous):
323
+ """
324
+ The noise predicition model function that is used for DPM-Solver.
325
+ """
326
+ if t_continuous.reshape((-1,)).shape[0] == 1:
327
+ t_continuous = t_continuous.expand((x.shape[0]))
328
+ if guidance_type == "uncond":
329
+ return noise_pred_fn(x, t_continuous)
330
+ elif guidance_type == "classifier":
331
+ assert classifier_fn is not None
332
+ t_input = get_model_input_time(t_continuous)
333
+ cond_grad = cond_grad_fn(x, t_input)
334
+ sigma_t = noise_schedule.marginal_std(t_continuous)
335
+ noise = noise_pred_fn(x, t_continuous)
336
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
337
+ elif guidance_type == "classifier-free":
338
+ if guidance_scale == 1. or unconditional_condition is None:
339
+ return noise_pred_fn(x, t_continuous, cond=condition)
340
+ else:
341
+ x_in = torch.cat([x] * 2)
342
+ t_in = torch.cat([t_continuous] * 2)
343
+ c_in = torch.cat([unconditional_condition, condition])
344
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
345
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
346
+
347
+ assert model_type in ["noise", "x_start", "v"]
348
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
349
+ return model_fn
350
+
351
+
352
+ class UniPC:
353
+ def __init__(
354
+ self,
355
+ model_fn,
356
+ noise_schedule,
357
+ predict_x0=True,
358
+ thresholding=False,
359
+ max_val=1.,
360
+ variant='bh1',
361
+ noise_mask=None,
362
+ masked_image=None,
363
+ noise=None,
364
+ ):
365
+ """Construct a UniPC.
366
+
367
+ We support both data_prediction and noise_prediction.
368
+ """
369
+ self.model = model_fn
370
+ self.noise_schedule = noise_schedule
371
+ self.variant = variant
372
+ self.predict_x0 = predict_x0
373
+ self.thresholding = thresholding
374
+ self.max_val = max_val
375
+ self.noise_mask = noise_mask
376
+ self.masked_image = masked_image
377
+ self.noise = noise
378
+
379
+ def dynamic_thresholding_fn(self, x0, t=None):
380
+ """
381
+ The dynamic thresholding method.
382
+ """
383
+ dims = x0.dim()
384
+ p = self.dynamic_thresholding_ratio
385
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
386
+ s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
387
+ x0 = torch.clamp(x0, -s, s) / s
388
+ return x0
389
+
390
+ def noise_prediction_fn(self, x, t):
391
+ """
392
+ Return the noise prediction model.
393
+ """
394
+ if self.noise_mask is not None:
395
+ return self.model(x, t) * self.noise_mask
396
+ else:
397
+ return self.model(x, t)
398
+
399
+ def data_prediction_fn(self, x, t):
400
+ """
401
+ Return the data prediction model (with thresholding).
402
+ """
403
+ noise = self.noise_prediction_fn(x, t)
404
+ dims = x.dim()
405
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
406
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
407
+ if self.thresholding:
408
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
409
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
410
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
411
+ x0 = torch.clamp(x0, -s, s) / s
412
+ if self.noise_mask is not None:
413
+ x0 = x0 * self.noise_mask + (1. - self.noise_mask) * self.masked_image
414
+ return x0
415
+
416
+ def model_fn(self, x, t):
417
+ """
418
+ Convert the model to the noise prediction model or the data prediction model.
419
+ """
420
+ if self.predict_x0:
421
+ return self.data_prediction_fn(x, t)
422
+ else:
423
+ return self.noise_prediction_fn(x, t)
424
+
425
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
426
+ """Compute the intermediate time steps for sampling.
427
+ """
428
+ if skip_type == 'logSNR':
429
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
430
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
431
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
432
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
433
+ elif skip_type == 'time_uniform':
434
+ return torch.linspace(t_T, t_0, N + 1).to(device)
435
+ elif skip_type == 'time_quadratic':
436
+ t_order = 2
437
+ t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
438
+ return t
439
+ else:
440
+ raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
441
+
442
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
443
+ """
444
+ Get the order of each step for sampling by the singlestep DPM-Solver.
445
+ """
446
+ if order == 3:
447
+ K = steps // 3 + 1
448
+ if steps % 3 == 0:
449
+ orders = [3,] * (K - 2) + [2, 1]
450
+ elif steps % 3 == 1:
451
+ orders = [3,] * (K - 1) + [1]
452
+ else:
453
+ orders = [3,] * (K - 1) + [2]
454
+ elif order == 2:
455
+ if steps % 2 == 0:
456
+ K = steps // 2
457
+ orders = [2,] * K
458
+ else:
459
+ K = steps // 2 + 1
460
+ orders = [2,] * (K - 1) + [1]
461
+ elif order == 1:
462
+ K = steps
463
+ orders = [1,] * steps
464
+ else:
465
+ raise ValueError("'order' must be '1' or '2' or '3'.")
466
+ if skip_type == 'logSNR':
467
+ # To reproduce the results in DPM-Solver paper
468
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
469
+ else:
470
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
471
+ return timesteps_outer, orders
472
+
473
+ def denoise_to_zero_fn(self, x, s):
474
+ """
475
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
476
+ """
477
+ return self.data_prediction_fn(x, s)
478
+
479
+ def multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
480
+ if len(t.shape) == 0:
481
+ t = t.view(-1)
482
+ if 'bh' in self.variant:
483
+ return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
484
+ else:
485
+ assert self.variant == 'vary_coeff'
486
+ return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
487
+
488
+ def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
489
+ print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
490
+ ns = self.noise_schedule
491
+ assert order <= len(model_prev_list)
492
+
493
+ # first compute rks
494
+ t_prev_0 = t_prev_list[-1]
495
+ lambda_prev_0 = ns.marginal_lambda(t_prev_0)
496
+ lambda_t = ns.marginal_lambda(t)
497
+ model_prev_0 = model_prev_list[-1]
498
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
499
+ log_alpha_t = ns.marginal_log_mean_coeff(t)
500
+ alpha_t = torch.exp(log_alpha_t)
501
+
502
+ h = lambda_t - lambda_prev_0
503
+
504
+ rks = []
505
+ D1s = []
506
+ for i in range(1, order):
507
+ t_prev_i = t_prev_list[-(i + 1)]
508
+ model_prev_i = model_prev_list[-(i + 1)]
509
+ lambda_prev_i = ns.marginal_lambda(t_prev_i)
510
+ rk = (lambda_prev_i - lambda_prev_0) / h
511
+ rks.append(rk)
512
+ D1s.append((model_prev_i - model_prev_0) / rk)
513
+
514
+ rks.append(1.)
515
+ rks = torch.tensor(rks, device=x.device)
516
+
517
+ K = len(rks)
518
+ # build C matrix
519
+ C = []
520
+
521
+ col = torch.ones_like(rks)
522
+ for k in range(1, K + 1):
523
+ C.append(col)
524
+ col = col * rks / (k + 1)
525
+ C = torch.stack(C, dim=1)
526
+
527
+ if len(D1s) > 0:
528
+ D1s = torch.stack(D1s, dim=1) # (B, K)
529
+ C_inv_p = torch.linalg.inv(C[:-1, :-1])
530
+ A_p = C_inv_p
531
+
532
+ if use_corrector:
533
+ print('using corrector')
534
+ C_inv = torch.linalg.inv(C)
535
+ A_c = C_inv
536
+
537
+ hh = -h if self.predict_x0 else h
538
+ h_phi_1 = torch.expm1(hh)
539
+ h_phi_ks = []
540
+ factorial_k = 1
541
+ h_phi_k = h_phi_1
542
+ for k in range(1, K + 2):
543
+ h_phi_ks.append(h_phi_k)
544
+ h_phi_k = h_phi_k / hh - 1 / factorial_k
545
+ factorial_k *= (k + 1)
546
+
547
+ model_t = None
548
+ if self.predict_x0:
549
+ x_t_ = (
550
+ sigma_t / sigma_prev_0 * x
551
+ - alpha_t * h_phi_1 * model_prev_0
552
+ )
553
+ # now predictor
554
+ x_t = x_t_
555
+ if len(D1s) > 0:
556
+ # compute the residuals for predictor
557
+ for k in range(K - 1):
558
+ x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
559
+ # now corrector
560
+ if use_corrector:
561
+ model_t = self.model_fn(x_t, t)
562
+ D1_t = (model_t - model_prev_0)
563
+ x_t = x_t_
564
+ k = 0
565
+ for k in range(K - 1):
566
+ x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
567
+ x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
568
+ else:
569
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
570
+ x_t_ = (
571
+ (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
572
+ - (sigma_t * h_phi_1) * model_prev_0
573
+ )
574
+ # now predictor
575
+ x_t = x_t_
576
+ if len(D1s) > 0:
577
+ # compute the residuals for predictor
578
+ for k in range(K - 1):
579
+ x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
580
+ # now corrector
581
+ if use_corrector:
582
+ model_t = self.model_fn(x_t, t)
583
+ D1_t = (model_t - model_prev_0)
584
+ x_t = x_t_
585
+ k = 0
586
+ for k in range(K - 1):
587
+ x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
588
+ x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
589
+ return x_t, model_t
590
+
591
+ def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
592
+ # print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
593
+ ns = self.noise_schedule
594
+ assert order <= len(model_prev_list)
595
+ dims = x.dim()
596
+
597
+ # first compute rks
598
+ t_prev_0 = t_prev_list[-1]
599
+ lambda_prev_0 = ns.marginal_lambda(t_prev_0)
600
+ lambda_t = ns.marginal_lambda(t)
601
+ model_prev_0 = model_prev_list[-1]
602
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
603
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
604
+ alpha_t = torch.exp(log_alpha_t)
605
+
606
+ h = lambda_t - lambda_prev_0
607
+
608
+ rks = []
609
+ D1s = []
610
+ for i in range(1, order):
611
+ t_prev_i = t_prev_list[-(i + 1)]
612
+ model_prev_i = model_prev_list[-(i + 1)]
613
+ lambda_prev_i = ns.marginal_lambda(t_prev_i)
614
+ rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
615
+ rks.append(rk)
616
+ D1s.append((model_prev_i - model_prev_0) / rk)
617
+
618
+ rks.append(1.)
619
+ rks = torch.tensor(rks, device=x.device)
620
+
621
+ R = []
622
+ b = []
623
+
624
+ hh = -h[0] if self.predict_x0 else h[0]
625
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
626
+ h_phi_k = h_phi_1 / hh - 1
627
+
628
+ factorial_i = 1
629
+
630
+ if self.variant == 'bh1':
631
+ B_h = hh
632
+ elif self.variant == 'bh2':
633
+ B_h = torch.expm1(hh)
634
+ else:
635
+ raise NotImplementedError()
636
+
637
+ for i in range(1, order + 1):
638
+ R.append(torch.pow(rks, i - 1))
639
+ b.append(h_phi_k * factorial_i / B_h)
640
+ factorial_i *= (i + 1)
641
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
642
+
643
+ R = torch.stack(R)
644
+ b = torch.tensor(b, device=x.device)
645
+
646
+ # now predictor
647
+ use_predictor = len(D1s) > 0 and x_t is None
648
+ if len(D1s) > 0:
649
+ D1s = torch.stack(D1s, dim=1) # (B, K)
650
+ if x_t is None:
651
+ # for order 2, we use a simplified version
652
+ if order == 2:
653
+ rhos_p = torch.tensor([0.5], device=b.device)
654
+ else:
655
+ rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
656
+ else:
657
+ D1s = None
658
+
659
+ if use_corrector:
660
+ # print('using corrector')
661
+ # for order 1, we use a simplified version
662
+ if order == 1:
663
+ rhos_c = torch.tensor([0.5], device=b.device)
664
+ else:
665
+ rhos_c = torch.linalg.solve(R, b)
666
+
667
+ model_t = None
668
+ if self.predict_x0:
669
+ x_t_ = (
670
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
671
+ - expand_dims(alpha_t * h_phi_1, dims)* model_prev_0
672
+ )
673
+
674
+ if x_t is None:
675
+ if use_predictor:
676
+ pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
677
+ else:
678
+ pred_res = 0
679
+ x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * pred_res
680
+
681
+ if use_corrector:
682
+ model_t = self.model_fn(x_t, t)
683
+ if D1s is not None:
684
+ corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
685
+ else:
686
+ corr_res = 0
687
+ D1_t = (model_t - model_prev_0)
688
+ x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
689
+ else:
690
+ x_t_ = (
691
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dimss) * x
692
+ - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
693
+ )
694
+ if x_t is None:
695
+ if use_predictor:
696
+ pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
697
+ else:
698
+ pred_res = 0
699
+ x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * pred_res
700
+
701
+ if use_corrector:
702
+ model_t = self.model_fn(x_t, t)
703
+ if D1s is not None:
704
+ corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
705
+ else:
706
+ corr_res = 0
707
+ D1_t = (model_t - model_prev_0)
708
+ x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
709
+ return x_t, model_t
710
+
711
+
712
+ def sample(self, x, timesteps, t_start=None, t_end=None, order=3, skip_type='time_uniform',
713
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
714
+ atol=0.0078, rtol=0.05, corrector=False, callback=None, disable_pbar=False
715
+ ):
716
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
717
+ t_T = self.noise_schedule.T if t_start is None else t_start
718
+ device = x.device
719
+ steps = len(timesteps) - 1
720
+ if method == 'multistep':
721
+ assert steps >= order
722
+ # timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
723
+ assert timesteps.shape[0] - 1 == steps
724
+ # with torch.no_grad():
725
+ for step_index in trange(steps, disable=disable_pbar):
726
+ if self.noise_mask is not None:
727
+ 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]))
728
+ if step_index == 0:
729
+ vec_t = timesteps[0].expand((x.shape[0]))
730
+ model_prev_list = [self.model_fn(x, vec_t)]
731
+ t_prev_list = [vec_t]
732
+ elif step_index < order:
733
+ init_order = step_index
734
+ # Init the first `order` values by lower order multistep DPM-Solver.
735
+ # for init_order in range(1, order):
736
+ vec_t = timesteps[init_order].expand(x.shape[0])
737
+ x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
738
+ if model_x is None:
739
+ model_x = self.model_fn(x, vec_t)
740
+ model_prev_list.append(model_x)
741
+ t_prev_list.append(vec_t)
742
+ else:
743
+ extra_final_step = 0
744
+ if step_index == (steps - 1):
745
+ extra_final_step = 1
746
+ for step in range(step_index, step_index + 1 + extra_final_step):
747
+ vec_t = timesteps[step].expand(x.shape[0])
748
+ if lower_order_final:
749
+ step_order = min(order, steps + 1 - step)
750
+ else:
751
+ step_order = order
752
+ # print('this step order:', step_order)
753
+ if step == steps:
754
+ # print('do not run corrector at the last step')
755
+ use_corrector = False
756
+ else:
757
+ use_corrector = True
758
+ x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
759
+ for i in range(order - 1):
760
+ t_prev_list[i] = t_prev_list[i + 1]
761
+ model_prev_list[i] = model_prev_list[i + 1]
762
+ t_prev_list[-1] = vec_t
763
+ # We do not need to evaluate the final model value.
764
+ if step < steps:
765
+ if model_x is None:
766
+ model_x = self.model_fn(x, vec_t)
767
+ model_prev_list[-1] = model_x
768
+ if callback is not None:
769
+ callback(step_index, model_prev_list[-1], x, steps)
770
+ else:
771
+ raise NotImplementedError()
772
+ if denoise_to_zero:
773
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
774
+ return x
775
+
776
+
777
+ #############################################################
778
+ # other utility functions
779
+ #############################################################
780
+
781
+ def interpolate_fn(x, xp, yp):
782
+ """
783
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
784
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
785
+ 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.)
786
+
787
+ Args:
788
+ 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).
789
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
790
+ yp: PyTorch tensor with shape [C, K].
791
+ Returns:
792
+ The function values f(x), with shape [N, C].
793
+ """
794
+ N, K = x.shape[0], xp.shape[1]
795
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
796
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
797
+ x_idx = torch.argmin(x_indices, dim=2)
798
+ cand_start_idx = x_idx - 1
799
+ start_idx = torch.where(
800
+ torch.eq(x_idx, 0),
801
+ torch.tensor(1, device=x.device),
802
+ torch.where(
803
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
804
+ ),
805
+ )
806
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
807
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
808
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
809
+ start_idx2 = torch.where(
810
+ torch.eq(x_idx, 0),
811
+ torch.tensor(0, device=x.device),
812
+ torch.where(
813
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
814
+ ),
815
+ )
816
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
817
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
818
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
819
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
820
+ return cand
821
+
822
+
823
+ def expand_dims(v, dims):
824
+ """
825
+ Expand the tensor `v` to the dim `dims`.
826
+
827
+ Args:
828
+ `v`: a PyTorch tensor with shape [N].
829
+ `dim`: a `int`.
830
+ Returns:
831
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
832
+ """
833
+ return v[(...,) + (None,)*(dims - 1)]
834
+
835
+
836
+
837
+ def sample_unipc(model, noise, image, sigmas, sampling_function, max_denoise, extra_args=None, callback=None, disable=False, noise_mask=None, variant='bh1'):
838
+ to_zero = False
839
+ if sigmas[-1] == 0:
840
+ timesteps = torch.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
841
+ to_zero = True
842
+ else:
843
+ timesteps = sigmas.clone()
844
+
845
+ alphas_cumprod = model.inner_model.alphas_cumprod
846
+
847
+ for s in range(timesteps.shape[0]):
848
+ timesteps[s] = (model.sigma_to_discrete_timestep(timesteps[s]) / 1000) + (1 / len(alphas_cumprod))
849
+
850
+ ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
851
+
852
+ if image is not None:
853
+ img = image * ns.marginal_alpha(timesteps[0])
854
+ if max_denoise:
855
+ noise_mult = 1.0
856
+ else:
857
+ noise_mult = ns.marginal_std(timesteps[0])
858
+ img += noise * noise_mult
859
+ else:
860
+ img = noise
861
+
862
+ if to_zero:
863
+ timesteps[-1] = (1 / len(alphas_cumprod))
864
+
865
+ device = noise.device
866
+
867
+
868
+ model_type = "noise"
869
+
870
+ model_fn = model_wrapper(
871
+ model.predict_eps_discrete_timestep,
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, callback=callback, disable_pbar=disable)
881
+ if not to_zero:
882
+ x /= ns.marginal_alpha(timesteps[-1])
883
+ return x
comfy/gligen.py ADDED
@@ -0,0 +1,341 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn, einsum
3
+ from .ldm.modules.attention import CrossAttention
4
+ from inspect import isfunction
5
+
6
+
7
+ def exists(val):
8
+ return val is not None
9
+
10
+
11
+ def uniq(arr):
12
+ return{el: True for el in arr}.keys()
13
+
14
+
15
+ def default(val, d):
16
+ if exists(val):
17
+ return val
18
+ return d() if isfunction(d) else d
19
+
20
+
21
+ # feedforward
22
+ class GEGLU(nn.Module):
23
+ def __init__(self, dim_in, dim_out):
24
+ super().__init__()
25
+ self.proj = nn.Linear(dim_in, dim_out * 2)
26
+
27
+ def forward(self, x):
28
+ x, gate = self.proj(x).chunk(2, dim=-1)
29
+ return x * torch.nn.functional.gelu(gate)
30
+
31
+
32
+ class FeedForward(nn.Module):
33
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
34
+ super().__init__()
35
+ inner_dim = int(dim * mult)
36
+ dim_out = default(dim_out, dim)
37
+ project_in = nn.Sequential(
38
+ nn.Linear(dim, inner_dim),
39
+ nn.GELU()
40
+ ) if not glu else GEGLU(dim, inner_dim)
41
+
42
+ self.net = nn.Sequential(
43
+ project_in,
44
+ nn.Dropout(dropout),
45
+ nn.Linear(inner_dim, dim_out)
46
+ )
47
+
48
+ def forward(self, x):
49
+ return self.net(x)
50
+
51
+
52
+ class GatedCrossAttentionDense(nn.Module):
53
+ def __init__(self, query_dim, context_dim, n_heads, d_head):
54
+ super().__init__()
55
+
56
+ self.attn = CrossAttention(
57
+ query_dim=query_dim,
58
+ context_dim=context_dim,
59
+ heads=n_heads,
60
+ dim_head=d_head)
61
+ self.ff = FeedForward(query_dim, glu=True)
62
+
63
+ self.norm1 = nn.LayerNorm(query_dim)
64
+ self.norm2 = nn.LayerNorm(query_dim)
65
+
66
+ self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
67
+ self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
68
+
69
+ # this can be useful: we can externally change magnitude of tanh(alpha)
70
+ # for example, when it is set to 0, then the entire model is same as
71
+ # original one
72
+ self.scale = 1
73
+
74
+ def forward(self, x, objs):
75
+
76
+ x = x + self.scale * \
77
+ torch.tanh(self.alpha_attn) * self.attn(self.norm1(x), objs, objs)
78
+ x = x + self.scale * \
79
+ torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
80
+
81
+ return x
82
+
83
+
84
+ class GatedSelfAttentionDense(nn.Module):
85
+ def __init__(self, query_dim, context_dim, n_heads, d_head):
86
+ super().__init__()
87
+
88
+ # we need a linear projection since we need cat visual feature and obj
89
+ # feature
90
+ self.linear = nn.Linear(context_dim, query_dim)
91
+
92
+ self.attn = CrossAttention(
93
+ query_dim=query_dim,
94
+ context_dim=query_dim,
95
+ heads=n_heads,
96
+ dim_head=d_head)
97
+ self.ff = FeedForward(query_dim, glu=True)
98
+
99
+ self.norm1 = nn.LayerNorm(query_dim)
100
+ self.norm2 = nn.LayerNorm(query_dim)
101
+
102
+ self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
103
+ self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
104
+
105
+ # this can be useful: we can externally change magnitude of tanh(alpha)
106
+ # for example, when it is set to 0, then the entire model is same as
107
+ # original one
108
+ self.scale = 1
109
+
110
+ def forward(self, x, objs):
111
+
112
+ N_visual = x.shape[1]
113
+ objs = self.linear(objs)
114
+
115
+ x = x + self.scale * torch.tanh(self.alpha_attn) * self.attn(
116
+ self.norm1(torch.cat([x, objs], dim=1)))[:, 0:N_visual, :]
117
+ x = x + self.scale * \
118
+ torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
119
+
120
+ return x
121
+
122
+
123
+ class GatedSelfAttentionDense2(nn.Module):
124
+ def __init__(self, query_dim, context_dim, n_heads, d_head):
125
+ super().__init__()
126
+
127
+ # we need a linear projection since we need cat visual feature and obj
128
+ # feature
129
+ self.linear = nn.Linear(context_dim, query_dim)
130
+
131
+ self.attn = CrossAttention(
132
+ query_dim=query_dim, context_dim=query_dim, dim_head=d_head)
133
+ self.ff = FeedForward(query_dim, glu=True)
134
+
135
+ self.norm1 = nn.LayerNorm(query_dim)
136
+ self.norm2 = nn.LayerNorm(query_dim)
137
+
138
+ self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
139
+ self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
140
+
141
+ # this can be useful: we can externally change magnitude of tanh(alpha)
142
+ # for example, when it is set to 0, then the entire model is same as
143
+ # original one
144
+ self.scale = 1
145
+
146
+ def forward(self, x, objs):
147
+
148
+ B, N_visual, _ = x.shape
149
+ B, N_ground, _ = objs.shape
150
+
151
+ objs = self.linear(objs)
152
+
153
+ # sanity check
154
+ size_v = math.sqrt(N_visual)
155
+ size_g = math.sqrt(N_ground)
156
+ assert int(size_v) == size_v, "Visual tokens must be square rootable"
157
+ assert int(size_g) == size_g, "Grounding tokens must be square rootable"
158
+ size_v = int(size_v)
159
+ size_g = int(size_g)
160
+
161
+ # select grounding token and resize it to visual token size as residual
162
+ out = self.attn(self.norm1(torch.cat([x, objs], dim=1)))[
163
+ :, N_visual:, :]
164
+ out = out.permute(0, 2, 1).reshape(B, -1, size_g, size_g)
165
+ out = torch.nn.functional.interpolate(
166
+ out, (size_v, size_v), mode='bicubic')
167
+ residual = out.reshape(B, -1, N_visual).permute(0, 2, 1)
168
+
169
+ # add residual to visual feature
170
+ x = x + self.scale * torch.tanh(self.alpha_attn) * residual
171
+ x = x + self.scale * \
172
+ torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
173
+
174
+ return x
175
+
176
+
177
+ class FourierEmbedder():
178
+ def __init__(self, num_freqs=64, temperature=100):
179
+
180
+ self.num_freqs = num_freqs
181
+ self.temperature = temperature
182
+ self.freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
183
+
184
+ @torch.no_grad()
185
+ def __call__(self, x, cat_dim=-1):
186
+ "x: arbitrary shape of tensor. dim: cat dim"
187
+ out = []
188
+ for freq in self.freq_bands:
189
+ out.append(torch.sin(freq * x))
190
+ out.append(torch.cos(freq * x))
191
+ return torch.cat(out, cat_dim)
192
+
193
+
194
+ class PositionNet(nn.Module):
195
+ def __init__(self, in_dim, out_dim, fourier_freqs=8):
196
+ super().__init__()
197
+ self.in_dim = in_dim
198
+ self.out_dim = out_dim
199
+
200
+ self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
201
+ self.position_dim = fourier_freqs * 2 * 4 # 2 is sin&cos, 4 is xyxy
202
+
203
+ self.linears = nn.Sequential(
204
+ nn.Linear(self.in_dim + self.position_dim, 512),
205
+ nn.SiLU(),
206
+ nn.Linear(512, 512),
207
+ nn.SiLU(),
208
+ nn.Linear(512, out_dim),
209
+ )
210
+
211
+ self.null_positive_feature = torch.nn.Parameter(
212
+ torch.zeros([self.in_dim]))
213
+ self.null_position_feature = torch.nn.Parameter(
214
+ torch.zeros([self.position_dim]))
215
+
216
+ def forward(self, boxes, masks, positive_embeddings):
217
+ B, N, _ = boxes.shape
218
+ dtype = self.linears[0].weight.dtype
219
+ masks = masks.unsqueeze(-1).to(dtype)
220
+ positive_embeddings = positive_embeddings.to(dtype)
221
+
222
+ # embedding position (it may includes padding as placeholder)
223
+ xyxy_embedding = self.fourier_embedder(boxes.to(dtype)) # B*N*4 --> B*N*C
224
+
225
+ # learnable null embedding
226
+ positive_null = self.null_positive_feature.view(1, 1, -1)
227
+ xyxy_null = self.null_position_feature.view(1, 1, -1)
228
+
229
+ # replace padding with learnable null embedding
230
+ positive_embeddings = positive_embeddings * \
231
+ masks + (1 - masks) * positive_null
232
+ xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null
233
+
234
+ objs = self.linears(
235
+ torch.cat([positive_embeddings, xyxy_embedding], dim=-1))
236
+ assert objs.shape == torch.Size([B, N, self.out_dim])
237
+ return objs
238
+
239
+
240
+ class Gligen(nn.Module):
241
+ def __init__(self, modules, position_net, key_dim):
242
+ super().__init__()
243
+ self.module_list = nn.ModuleList(modules)
244
+ self.position_net = position_net
245
+ self.key_dim = key_dim
246
+ self.max_objs = 30
247
+ self.current_device = torch.device("cpu")
248
+
249
+ def _set_position(self, boxes, masks, positive_embeddings):
250
+ objs = self.position_net(boxes, masks, positive_embeddings)
251
+ def func(x, extra_options):
252
+ key = extra_options["transformer_index"]
253
+ module = self.module_list[key]
254
+ return module(x, objs)
255
+ return func
256
+
257
+ def set_position(self, latent_image_shape, position_params, device):
258
+ batch, c, h, w = latent_image_shape
259
+ masks = torch.zeros([self.max_objs], device="cpu")
260
+ boxes = []
261
+ positive_embeddings = []
262
+ for p in position_params:
263
+ x1 = (p[4]) / w
264
+ y1 = (p[3]) / h
265
+ x2 = (p[4] + p[2]) / w
266
+ y2 = (p[3] + p[1]) / h
267
+ masks[len(boxes)] = 1.0
268
+ boxes += [torch.tensor((x1, y1, x2, y2)).unsqueeze(0)]
269
+ positive_embeddings += [p[0]]
270
+ append_boxes = []
271
+ append_conds = []
272
+ if len(boxes) < self.max_objs:
273
+ append_boxes = [torch.zeros(
274
+ [self.max_objs - len(boxes), 4], device="cpu")]
275
+ append_conds = [torch.zeros(
276
+ [self.max_objs - len(boxes), self.key_dim], device="cpu")]
277
+
278
+ box_out = torch.cat(
279
+ boxes + append_boxes).unsqueeze(0).repeat(batch, 1, 1)
280
+ masks = masks.unsqueeze(0).repeat(batch, 1)
281
+ conds = torch.cat(positive_embeddings +
282
+ append_conds).unsqueeze(0).repeat(batch, 1, 1)
283
+ return self._set_position(
284
+ box_out.to(device),
285
+ masks.to(device),
286
+ conds.to(device))
287
+
288
+ def set_empty(self, latent_image_shape, device):
289
+ batch, c, h, w = latent_image_shape
290
+ masks = torch.zeros([self.max_objs], device="cpu").repeat(batch, 1)
291
+ box_out = torch.zeros([self.max_objs, 4],
292
+ device="cpu").repeat(batch, 1, 1)
293
+ conds = torch.zeros([self.max_objs, self.key_dim],
294
+ device="cpu").repeat(batch, 1, 1)
295
+ return self._set_position(
296
+ box_out.to(device),
297
+ masks.to(device),
298
+ conds.to(device))
299
+
300
+
301
+ def load_gligen(sd):
302
+ sd_k = sd.keys()
303
+ output_list = []
304
+ key_dim = 768
305
+ for a in ["input_blocks", "middle_block", "output_blocks"]:
306
+ for b in range(20):
307
+ k_temp = filter(lambda k: "{}.{}.".format(a, b)
308
+ in k and ".fuser." in k, sd_k)
309
+ k_temp = map(lambda k: (k, k.split(".fuser.")[-1]), k_temp)
310
+
311
+ n_sd = {}
312
+ for k in k_temp:
313
+ n_sd[k[1]] = sd[k[0]]
314
+ if len(n_sd) > 0:
315
+ query_dim = n_sd["linear.weight"].shape[0]
316
+ key_dim = n_sd["linear.weight"].shape[1]
317
+
318
+ if key_dim == 768: # SD1.x
319
+ n_heads = 8
320
+ d_head = query_dim // n_heads
321
+ else:
322
+ d_head = 64
323
+ n_heads = query_dim // d_head
324
+
325
+ gated = GatedSelfAttentionDense(
326
+ query_dim, key_dim, n_heads, d_head)
327
+ gated.load_state_dict(n_sd, strict=False)
328
+ output_list.append(gated)
329
+
330
+ if "position_net.null_positive_feature" in sd_k:
331
+ in_dim = sd["position_net.null_positive_feature"].shape[0]
332
+ out_dim = sd["position_net.linears.4.weight"].shape[0]
333
+
334
+ class WeightsLoader(torch.nn.Module):
335
+ pass
336
+ w = WeightsLoader()
337
+ w.position_net = PositionNet(in_dim, out_dim)
338
+ w.load_state_dict(sd, strict=False)
339
+
340
+ gligen = Gligen(output_list, w.position_net, key_dim)
341
+ return gligen
comfy/k_diffusion/external.py ADDED
@@ -0,0 +1,190 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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_discrete_timestep(self, sigma):
67
+ log_sigma = sigma.log()
68
+ dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
69
+ return dists.abs().argmin(dim=0).view(sigma.shape)
70
+
71
+ def sigma_to_t(self, sigma, quantize=None):
72
+ quantize = self.quantize if quantize is None else quantize
73
+ if quantize:
74
+ return self.sigma_to_discrete_timestep(sigma)
75
+ log_sigma = sigma.log()
76
+ dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
77
+ low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2)
78
+ high_idx = low_idx + 1
79
+ low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx]
80
+ w = (low - log_sigma) / (low - high)
81
+ w = w.clamp(0, 1)
82
+ t = (1 - w) * low_idx + w * high_idx
83
+ return t.view(sigma.shape)
84
+
85
+ def t_to_sigma(self, t):
86
+ t = t.float()
87
+ low_idx = t.floor().long()
88
+ high_idx = t.ceil().long()
89
+ w = t-low_idx if t.device.type == 'mps' else t.frac()
90
+ log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
91
+ return log_sigma.exp()
92
+
93
+ def predict_eps_discrete_timestep(self, input, t, **kwargs):
94
+ if t.dtype != torch.int64 and t.dtype != torch.int32:
95
+ t = t.round()
96
+ sigma = self.t_to_sigma(t)
97
+ input = input * ((utils.append_dims(sigma, input.ndim) ** 2 + 1.0) ** 0.5)
98
+ return (input - self(input, sigma, **kwargs)) / utils.append_dims(sigma, input.ndim)
99
+
100
+ class DiscreteEpsDDPMDenoiser(DiscreteSchedule):
101
+ """A wrapper for discrete schedule DDPM models that output eps (the predicted
102
+ noise)."""
103
+
104
+ def __init__(self, model, alphas_cumprod, quantize):
105
+ super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize)
106
+ self.inner_model = model
107
+ self.sigma_data = 1.
108
+
109
+ def get_scalings(self, sigma):
110
+ c_out = -sigma
111
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
112
+ return c_out, c_in
113
+
114
+ def get_eps(self, *args, **kwargs):
115
+ return self.inner_model(*args, **kwargs)
116
+
117
+ def loss(self, input, noise, sigma, **kwargs):
118
+ c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
119
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
120
+ eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs)
121
+ return (eps - noise).pow(2).flatten(1).mean(1)
122
+
123
+ def forward(self, input, sigma, **kwargs):
124
+ c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
125
+ eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
126
+ return input + eps * c_out
127
+
128
+
129
+ class OpenAIDenoiser(DiscreteEpsDDPMDenoiser):
130
+ """A wrapper for OpenAI diffusion models."""
131
+
132
+ def __init__(self, model, diffusion, quantize=False, has_learned_sigmas=True, device='cpu'):
133
+ alphas_cumprod = torch.tensor(diffusion.alphas_cumprod, device=device, dtype=torch.float32)
134
+ super().__init__(model, alphas_cumprod, quantize=quantize)
135
+ self.has_learned_sigmas = has_learned_sigmas
136
+
137
+ def get_eps(self, *args, **kwargs):
138
+ model_output = self.inner_model(*args, **kwargs)
139
+ if self.has_learned_sigmas:
140
+ return model_output.chunk(2, dim=1)[0]
141
+ return model_output
142
+
143
+
144
+ class CompVisDenoiser(DiscreteEpsDDPMDenoiser):
145
+ """A wrapper for CompVis diffusion models."""
146
+
147
+ def __init__(self, model, quantize=False, device='cpu'):
148
+ super().__init__(model, model.alphas_cumprod, quantize=quantize)
149
+
150
+ def get_eps(self, *args, **kwargs):
151
+ return self.inner_model.apply_model(*args, **kwargs)
152
+
153
+
154
+ class DiscreteVDDPMDenoiser(DiscreteSchedule):
155
+ """A wrapper for discrete schedule DDPM models that output v."""
156
+
157
+ def __init__(self, model, alphas_cumprod, quantize):
158
+ super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize)
159
+ self.inner_model = model
160
+ self.sigma_data = 1.
161
+
162
+ def get_scalings(self, sigma):
163
+ c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
164
+ c_out = -sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
165
+ c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
166
+ return c_skip, c_out, c_in
167
+
168
+ def get_v(self, *args, **kwargs):
169
+ return self.inner_model(*args, **kwargs)
170
+
171
+ def loss(self, input, noise, sigma, **kwargs):
172
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
173
+ noised_input = input + noise * utils.append_dims(sigma, input.ndim)
174
+ model_output = self.get_v(noised_input * c_in, self.sigma_to_t(sigma), **kwargs)
175
+ target = (input - c_skip * noised_input) / c_out
176
+ return (model_output - target).pow(2).flatten(1).mean(1)
177
+
178
+ def forward(self, input, sigma, **kwargs):
179
+ c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
180
+ return self.get_v(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip
181
+
182
+
183
+ class CompVisVDenoiser(DiscreteVDDPMDenoiser):
184
+ """A wrapper for CompVis diffusion models that output v."""
185
+
186
+ def __init__(self, model, quantize=False, device='cpu'):
187
+ super().__init__(model, model.alphas_cumprod, quantize=quantize)
188
+
189
+ def get_v(self, x, t, cond, **kwargs):
190
+ return self.inner_model.apply_model(x, t, cond)
comfy/k_diffusion/sampling.py ADDED
@@ -0,0 +1,739 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ from scipy import integrate
4
+ import torch
5
+ from torch import nn
6
+ import torchsde
7
+ from tqdm.auto import trange, tqdm
8
+
9
+ from . import utils
10
+
11
+
12
+ def append_zero(x):
13
+ return torch.cat([x, x.new_zeros([1])])
14
+
15
+
16
+ def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
17
+ """Constructs the noise schedule of Karras et al. (2022)."""
18
+ ramp = torch.linspace(0, 1, n, device=device)
19
+ min_inv_rho = sigma_min ** (1 / rho)
20
+ max_inv_rho = sigma_max ** (1 / rho)
21
+ sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
22
+ return append_zero(sigmas).to(device)
23
+
24
+
25
+ def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
26
+ """Constructs an exponential noise schedule."""
27
+ sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
28
+ return append_zero(sigmas)
29
+
30
+
31
+ def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
32
+ """Constructs an polynomial in log sigma noise schedule."""
33
+ ramp = torch.linspace(1, 0, n, device=device) ** rho
34
+ sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
35
+ return append_zero(sigmas)
36
+
37
+
38
+ def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
39
+ """Constructs a continuous VP noise schedule."""
40
+ t = torch.linspace(1, eps_s, n, device=device)
41
+ sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
42
+ return append_zero(sigmas)
43
+
44
+
45
+ def to_d(x, sigma, denoised):
46
+ """Converts a denoiser output to a Karras ODE derivative."""
47
+ return (x - denoised) / utils.append_dims(sigma, x.ndim)
48
+
49
+
50
+ def get_ancestral_step(sigma_from, sigma_to, eta=1.):
51
+ """Calculates the noise level (sigma_down) to step down to and the amount
52
+ of noise to add (sigma_up) when doing an ancestral sampling step."""
53
+ if not eta:
54
+ return sigma_to, 0.
55
+ sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
56
+ sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
57
+ return sigma_down, sigma_up
58
+
59
+
60
+ def default_noise_sampler(x):
61
+ return lambda sigma, sigma_next: torch.randn_like(x)
62
+
63
+
64
+ class BatchedBrownianTree:
65
+ """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
66
+
67
+ def __init__(self, x, t0, t1, seed=None, **kwargs):
68
+ self.cpu_tree = True
69
+ if "cpu" in kwargs:
70
+ self.cpu_tree = kwargs.pop("cpu")
71
+ t0, t1, self.sign = self.sort(t0, t1)
72
+ w0 = kwargs.get('w0', torch.zeros_like(x))
73
+ if seed is None:
74
+ seed = torch.randint(0, 2 ** 63 - 1, []).item()
75
+ self.batched = True
76
+ try:
77
+ assert len(seed) == x.shape[0]
78
+ w0 = w0[0]
79
+ except TypeError:
80
+ seed = [seed]
81
+ self.batched = False
82
+ if self.cpu_tree:
83
+ self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
84
+ else:
85
+ self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
86
+
87
+ @staticmethod
88
+ def sort(a, b):
89
+ return (a, b, 1) if a < b else (b, a, -1)
90
+
91
+ def __call__(self, t0, t1):
92
+ t0, t1, sign = self.sort(t0, t1)
93
+ if self.cpu_tree:
94
+ w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
95
+ else:
96
+ w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
97
+
98
+ return w if self.batched else w[0]
99
+
100
+
101
+ class BrownianTreeNoiseSampler:
102
+ """A noise sampler backed by a torchsde.BrownianTree.
103
+
104
+ Args:
105
+ x (Tensor): The tensor whose shape, device and dtype to use to generate
106
+ random samples.
107
+ sigma_min (float): The low end of the valid interval.
108
+ sigma_max (float): The high end of the valid interval.
109
+ seed (int or List[int]): The random seed. If a list of seeds is
110
+ supplied instead of a single integer, then the noise sampler will
111
+ use one BrownianTree per batch item, each with its own seed.
112
+ transform (callable): A function that maps sigma to the sampler's
113
+ internal timestep.
114
+ """
115
+
116
+ def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
117
+ self.transform = transform
118
+ t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
119
+ self.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
120
+
121
+ def __call__(self, sigma, sigma_next):
122
+ t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
123
+ return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
124
+
125
+
126
+ @torch.no_grad()
127
+ 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.):
128
+ """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
129
+ extra_args = {} if extra_args is None else extra_args
130
+ s_in = x.new_ones([x.shape[0]])
131
+ for i in trange(len(sigmas) - 1, disable=disable):
132
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
133
+ sigma_hat = sigmas[i] * (gamma + 1)
134
+ if gamma > 0:
135
+ eps = torch.randn_like(x) * s_noise
136
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
137
+ denoised = model(x, sigma_hat * s_in, **extra_args)
138
+ d = to_d(x, sigma_hat, denoised)
139
+ if callback is not None:
140
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
141
+ dt = sigmas[i + 1] - sigma_hat
142
+ # Euler method
143
+ x = x + d * dt
144
+ return x
145
+
146
+
147
+ @torch.no_grad()
148
+ def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
149
+ """Ancestral sampling with Euler method steps."""
150
+ extra_args = {} if extra_args is None else extra_args
151
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
152
+ s_in = x.new_ones([x.shape[0]])
153
+ for i in trange(len(sigmas) - 1, disable=disable):
154
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
155
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
156
+ if callback is not None:
157
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
158
+ d = to_d(x, sigmas[i], denoised)
159
+ # Euler method
160
+ dt = sigma_down - sigmas[i]
161
+ x = x + d * dt
162
+ if sigmas[i + 1] > 0:
163
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
164
+ return x
165
+
166
+
167
+ @torch.no_grad()
168
+ 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.):
169
+ """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
170
+ extra_args = {} if extra_args is None else extra_args
171
+ s_in = x.new_ones([x.shape[0]])
172
+ for i in trange(len(sigmas) - 1, disable=disable):
173
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
174
+ sigma_hat = sigmas[i] * (gamma + 1)
175
+ if gamma > 0:
176
+ eps = torch.randn_like(x) * s_noise
177
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
178
+ denoised = model(x, sigma_hat * s_in, **extra_args)
179
+ d = to_d(x, sigma_hat, denoised)
180
+ if callback is not None:
181
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
182
+ dt = sigmas[i + 1] - sigma_hat
183
+ if sigmas[i + 1] == 0:
184
+ # Euler method
185
+ x = x + d * dt
186
+ else:
187
+ # Heun's method
188
+ x_2 = x + d * dt
189
+ denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
190
+ d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
191
+ d_prime = (d + d_2) / 2
192
+ x = x + d_prime * dt
193
+ return x
194
+
195
+
196
+ @torch.no_grad()
197
+ 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.):
198
+ """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
199
+ extra_args = {} if extra_args is None else extra_args
200
+ s_in = x.new_ones([x.shape[0]])
201
+ for i in trange(len(sigmas) - 1, disable=disable):
202
+ gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
203
+ sigma_hat = sigmas[i] * (gamma + 1)
204
+ if gamma > 0:
205
+ eps = torch.randn_like(x) * s_noise
206
+ x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
207
+ denoised = model(x, sigma_hat * s_in, **extra_args)
208
+ d = to_d(x, sigma_hat, denoised)
209
+ if callback is not None:
210
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
211
+ if sigmas[i + 1] == 0:
212
+ # Euler method
213
+ dt = sigmas[i + 1] - sigma_hat
214
+ x = x + d * dt
215
+ else:
216
+ # DPM-Solver-2
217
+ sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
218
+ dt_1 = sigma_mid - sigma_hat
219
+ dt_2 = sigmas[i + 1] - sigma_hat
220
+ x_2 = x + d * dt_1
221
+ denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
222
+ d_2 = to_d(x_2, sigma_mid, denoised_2)
223
+ x = x + d_2 * dt_2
224
+ return x
225
+
226
+
227
+ @torch.no_grad()
228
+ def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
229
+ """Ancestral sampling with DPM-Solver second-order steps."""
230
+ extra_args = {} if extra_args is None else extra_args
231
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
232
+ s_in = x.new_ones([x.shape[0]])
233
+ for i in trange(len(sigmas) - 1, disable=disable):
234
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
235
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
236
+ if callback is not None:
237
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
238
+ d = to_d(x, sigmas[i], denoised)
239
+ if sigma_down == 0:
240
+ # Euler method
241
+ dt = sigma_down - sigmas[i]
242
+ x = x + d * dt
243
+ else:
244
+ # DPM-Solver-2
245
+ sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
246
+ dt_1 = sigma_mid - sigmas[i]
247
+ dt_2 = sigma_down - sigmas[i]
248
+ x_2 = x + d * dt_1
249
+ denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
250
+ d_2 = to_d(x_2, sigma_mid, denoised_2)
251
+ x = x + d_2 * dt_2
252
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
253
+ return x
254
+
255
+
256
+ def linear_multistep_coeff(order, t, i, j):
257
+ if order - 1 > i:
258
+ raise ValueError(f'Order {order} too high for step {i}')
259
+ def fn(tau):
260
+ prod = 1.
261
+ for k in range(order):
262
+ if j == k:
263
+ continue
264
+ prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
265
+ return prod
266
+ return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
267
+
268
+
269
+ @torch.no_grad()
270
+ def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
271
+ extra_args = {} if extra_args is None else extra_args
272
+ s_in = x.new_ones([x.shape[0]])
273
+ sigmas_cpu = sigmas.detach().cpu().numpy()
274
+ ds = []
275
+ for i in trange(len(sigmas) - 1, disable=disable):
276
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
277
+ d = to_d(x, sigmas[i], denoised)
278
+ ds.append(d)
279
+ if len(ds) > order:
280
+ ds.pop(0)
281
+ if callback is not None:
282
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
283
+ cur_order = min(i + 1, order)
284
+ coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
285
+ x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
286
+ return x
287
+
288
+
289
+ class PIDStepSizeController:
290
+ """A PID controller for ODE adaptive step size control."""
291
+ def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
292
+ self.h = h
293
+ self.b1 = (pcoeff + icoeff + dcoeff) / order
294
+ self.b2 = -(pcoeff + 2 * dcoeff) / order
295
+ self.b3 = dcoeff / order
296
+ self.accept_safety = accept_safety
297
+ self.eps = eps
298
+ self.errs = []
299
+
300
+ def limiter(self, x):
301
+ return 1 + math.atan(x - 1)
302
+
303
+ def propose_step(self, error):
304
+ inv_error = 1 / (float(error) + self.eps)
305
+ if not self.errs:
306
+ self.errs = [inv_error, inv_error, inv_error]
307
+ self.errs[0] = inv_error
308
+ factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
309
+ factor = self.limiter(factor)
310
+ accept = factor >= self.accept_safety
311
+ if accept:
312
+ self.errs[2] = self.errs[1]
313
+ self.errs[1] = self.errs[0]
314
+ self.h *= factor
315
+ return accept
316
+
317
+
318
+ class DPMSolver(nn.Module):
319
+ """DPM-Solver. See https://arxiv.org/abs/2206.00927."""
320
+
321
+ def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None):
322
+ super().__init__()
323
+ self.model = model
324
+ self.extra_args = {} if extra_args is None else extra_args
325
+ self.eps_callback = eps_callback
326
+ self.info_callback = info_callback
327
+
328
+ def t(self, sigma):
329
+ return -sigma.log()
330
+
331
+ def sigma(self, t):
332
+ return t.neg().exp()
333
+
334
+ def eps(self, eps_cache, key, x, t, *args, **kwargs):
335
+ if key in eps_cache:
336
+ return eps_cache[key], eps_cache
337
+ sigma = self.sigma(t) * x.new_ones([x.shape[0]])
338
+ eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t)
339
+ if self.eps_callback is not None:
340
+ self.eps_callback()
341
+ return eps, {key: eps, **eps_cache}
342
+
343
+ def dpm_solver_1_step(self, x, t, t_next, eps_cache=None):
344
+ eps_cache = {} if eps_cache is None else eps_cache
345
+ h = t_next - t
346
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
347
+ x_1 = x - self.sigma(t_next) * h.expm1() * eps
348
+ return x_1, eps_cache
349
+
350
+ def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None):
351
+ eps_cache = {} if eps_cache is None else eps_cache
352
+ h = t_next - t
353
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
354
+ s1 = t + r1 * h
355
+ u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
356
+ eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
357
+ x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
358
+ return x_2, eps_cache
359
+
360
+ def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None):
361
+ eps_cache = {} if eps_cache is None else eps_cache
362
+ h = t_next - t
363
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
364
+ s1 = t + r1 * h
365
+ s2 = t + r2 * h
366
+ u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
367
+ eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
368
+ u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
369
+ eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2)
370
+ x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
371
+ return x_3, eps_cache
372
+
373
+ def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
374
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
375
+ if not t_end > t_start and eta:
376
+ raise ValueError('eta must be 0 for reverse sampling')
377
+
378
+ m = math.floor(nfe / 3) + 1
379
+ ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
380
+
381
+ if nfe % 3 == 0:
382
+ orders = [3] * (m - 2) + [2, 1]
383
+ else:
384
+ orders = [3] * (m - 1) + [nfe % 3]
385
+
386
+ for i in range(len(orders)):
387
+ eps_cache = {}
388
+ t, t_next = ts[i], ts[i + 1]
389
+ if eta:
390
+ sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
391
+ t_next_ = torch.minimum(t_end, self.t(sd))
392
+ su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
393
+ else:
394
+ t_next_, su = t_next, 0.
395
+
396
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
397
+ denoised = x - self.sigma(t) * eps
398
+ if self.info_callback is not None:
399
+ self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
400
+
401
+ if orders[i] == 1:
402
+ x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache)
403
+ elif orders[i] == 2:
404
+ x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)
405
+ else:
406
+ x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)
407
+
408
+ x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
409
+
410
+ return x
411
+
412
+ 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):
413
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
414
+ if order not in {2, 3}:
415
+ raise ValueError('order should be 2 or 3')
416
+ forward = t_end > t_start
417
+ if not forward and eta:
418
+ raise ValueError('eta must be 0 for reverse sampling')
419
+ h_init = abs(h_init) * (1 if forward else -1)
420
+ atol = torch.tensor(atol)
421
+ rtol = torch.tensor(rtol)
422
+ s = t_start
423
+ x_prev = x
424
+ accept = True
425
+ pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
426
+ info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
427
+
428
+ while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
429
+ eps_cache = {}
430
+ t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
431
+ if eta:
432
+ sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
433
+ t_ = torch.minimum(t_end, self.t(sd))
434
+ su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
435
+ else:
436
+ t_, su = t, 0.
437
+
438
+ eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
439
+ denoised = x - self.sigma(s) * eps
440
+
441
+ if order == 2:
442
+ x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
443
+ x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
444
+ else:
445
+ x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
446
+ x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
447
+ delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
448
+ error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
449
+ accept = pid.propose_step(error)
450
+ if accept:
451
+ x_prev = x_low
452
+ x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
453
+ s = t
454
+ info['n_accept'] += 1
455
+ else:
456
+ info['n_reject'] += 1
457
+ info['nfe'] += order
458
+ info['steps'] += 1
459
+
460
+ if self.info_callback is not None:
461
+ self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
462
+
463
+ return x, info
464
+
465
+
466
+ @torch.no_grad()
467
+ 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):
468
+ """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
469
+ if sigma_min <= 0 or sigma_max <= 0:
470
+ raise ValueError('sigma_min and sigma_max must not be 0')
471
+ with tqdm(total=n, disable=disable) as pbar:
472
+ dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
473
+ if callback is not None:
474
+ dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
475
+ 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)
476
+
477
+
478
+ @torch.no_grad()
479
+ 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):
480
+ """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
481
+ if sigma_min <= 0 or sigma_max <= 0:
482
+ raise ValueError('sigma_min and sigma_max must not be 0')
483
+ with tqdm(disable=disable) as pbar:
484
+ dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
485
+ if callback is not None:
486
+ dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
487
+ 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)
488
+ if return_info:
489
+ return x, info
490
+ return x
491
+
492
+
493
+ @torch.no_grad()
494
+ def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
495
+ """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
496
+ extra_args = {} if extra_args is None else extra_args
497
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
498
+ s_in = x.new_ones([x.shape[0]])
499
+ sigma_fn = lambda t: t.neg().exp()
500
+ t_fn = lambda sigma: sigma.log().neg()
501
+
502
+ for i in trange(len(sigmas) - 1, disable=disable):
503
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
504
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
505
+ if callback is not None:
506
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
507
+ if sigma_down == 0:
508
+ # Euler method
509
+ d = to_d(x, sigmas[i], denoised)
510
+ dt = sigma_down - sigmas[i]
511
+ x = x + d * dt
512
+ else:
513
+ # DPM-Solver++(2S)
514
+ t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
515
+ r = 1 / 2
516
+ h = t_next - t
517
+ s = t + r * h
518
+ x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
519
+ denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
520
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
521
+ # Noise addition
522
+ if sigmas[i + 1] > 0:
523
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
524
+ return x
525
+
526
+
527
+ @torch.no_grad()
528
+ 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):
529
+ """DPM-Solver++ (stochastic)."""
530
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
531
+ seed = extra_args.get("seed", None)
532
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
533
+ extra_args = {} if extra_args is None else extra_args
534
+ s_in = x.new_ones([x.shape[0]])
535
+ sigma_fn = lambda t: t.neg().exp()
536
+ t_fn = lambda sigma: sigma.log().neg()
537
+
538
+ for i in trange(len(sigmas) - 1, disable=disable):
539
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
540
+ if callback is not None:
541
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
542
+ if sigmas[i + 1] == 0:
543
+ # Euler method
544
+ d = to_d(x, sigmas[i], denoised)
545
+ dt = sigmas[i + 1] - sigmas[i]
546
+ x = x + d * dt
547
+ else:
548
+ # DPM-Solver++
549
+ t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
550
+ h = t_next - t
551
+ s = t + h * r
552
+ fac = 1 / (2 * r)
553
+
554
+ # Step 1
555
+ sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
556
+ s_ = t_fn(sd)
557
+ x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised
558
+ x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
559
+ denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
560
+
561
+ # Step 2
562
+ sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
563
+ t_next_ = t_fn(sd)
564
+ denoised_d = (1 - fac) * denoised + fac * denoised_2
565
+ x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d
566
+ x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
567
+ return x
568
+
569
+
570
+ @torch.no_grad()
571
+ def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
572
+ """DPM-Solver++(2M)."""
573
+ extra_args = {} if extra_args is None else extra_args
574
+ s_in = x.new_ones([x.shape[0]])
575
+ sigma_fn = lambda t: t.neg().exp()
576
+ t_fn = lambda sigma: sigma.log().neg()
577
+ old_denoised = None
578
+
579
+ for i in trange(len(sigmas) - 1, disable=disable):
580
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
581
+ if callback is not None:
582
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
583
+ t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
584
+ h = t_next - t
585
+ if old_denoised is None or sigmas[i + 1] == 0:
586
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
587
+ else:
588
+ h_last = t - t_fn(sigmas[i - 1])
589
+ r = h_last / h
590
+ denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
591
+ x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
592
+ old_denoised = denoised
593
+ return x
594
+
595
+ @torch.no_grad()
596
+ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
597
+ """DPM-Solver++(2M) SDE."""
598
+
599
+ if solver_type not in {'heun', 'midpoint'}:
600
+ raise ValueError('solver_type must be \'heun\' or \'midpoint\'')
601
+
602
+ seed = extra_args.get("seed", None)
603
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
604
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
605
+ extra_args = {} if extra_args is None else extra_args
606
+ s_in = x.new_ones([x.shape[0]])
607
+
608
+ old_denoised = None
609
+ h_last = None
610
+ h = None
611
+
612
+ for i in trange(len(sigmas) - 1, disable=disable):
613
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
614
+ if callback is not None:
615
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
616
+ if sigmas[i + 1] == 0:
617
+ # Denoising step
618
+ x = denoised
619
+ else:
620
+ # DPM-Solver++(2M) SDE
621
+ t, s = -sigmas[i].log(), -sigmas[i + 1].log()
622
+ h = s - t
623
+ eta_h = eta * h
624
+
625
+ x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
626
+
627
+ if old_denoised is not None:
628
+ r = h_last / h
629
+ if solver_type == 'heun':
630
+ x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
631
+ elif solver_type == 'midpoint':
632
+ x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
633
+
634
+ if eta:
635
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
636
+
637
+ old_denoised = denoised
638
+ h_last = h
639
+ return x
640
+
641
+ @torch.no_grad()
642
+ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
643
+ """DPM-Solver++(3M) SDE."""
644
+
645
+ seed = extra_args.get("seed", None)
646
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
647
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
648
+ extra_args = {} if extra_args is None else extra_args
649
+ s_in = x.new_ones([x.shape[0]])
650
+
651
+ denoised_1, denoised_2 = None, None
652
+ h, h_1, h_2 = None, None, None
653
+
654
+ for i in trange(len(sigmas) - 1, disable=disable):
655
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
656
+ if callback is not None:
657
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
658
+ if sigmas[i + 1] == 0:
659
+ # Denoising step
660
+ x = denoised
661
+ else:
662
+ t, s = -sigmas[i].log(), -sigmas[i + 1].log()
663
+ h = s - t
664
+ h_eta = h * (eta + 1)
665
+
666
+ x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised
667
+
668
+ if h_2 is not None:
669
+ r0 = h_1 / h
670
+ r1 = h_2 / h
671
+ d1_0 = (denoised - denoised_1) / r0
672
+ d1_1 = (denoised_1 - denoised_2) / r1
673
+ d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
674
+ d2 = (d1_0 - d1_1) / (r0 + r1)
675
+ phi_2 = h_eta.neg().expm1() / h_eta + 1
676
+ phi_3 = phi_2 / h_eta - 0.5
677
+ x = x + phi_2 * d1 - phi_3 * d2
678
+ elif h_1 is not None:
679
+ r = h_1 / h
680
+ d = (denoised - denoised_1) / r
681
+ phi_2 = h_eta.neg().expm1() / h_eta + 1
682
+ x = x + phi_2 * d
683
+
684
+ if eta:
685
+ x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
686
+
687
+ denoised_1, denoised_2 = denoised, denoised_1
688
+ h_1, h_2 = h, h_1
689
+ return x
690
+
691
+ @torch.no_grad()
692
+ def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
693
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
694
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
695
+ return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
696
+
697
+ @torch.no_grad()
698
+ def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
699
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
700
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
701
+ return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
702
+
703
+ @torch.no_grad()
704
+ def sample_dpmpp_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
705
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
706
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
707
+ return sample_dpmpp_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r)
708
+
709
+
710
+ def DDPMSampler_step(x, sigma, sigma_prev, noise, noise_sampler):
711
+ alpha_cumprod = 1 / ((sigma * sigma) + 1)
712
+ alpha_cumprod_prev = 1 / ((sigma_prev * sigma_prev) + 1)
713
+ alpha = (alpha_cumprod / alpha_cumprod_prev)
714
+
715
+ mu = (1.0 / alpha).sqrt() * (x - (1 - alpha) * noise / (1 - alpha_cumprod).sqrt())
716
+ if sigma_prev > 0:
717
+ mu += ((1 - alpha) * (1. - alpha_cumprod_prev) / (1. - alpha_cumprod)).sqrt() * noise_sampler(sigma, sigma_prev)
718
+ return mu
719
+
720
+
721
+ def generic_step_sampler(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, step_function=None):
722
+ extra_args = {} if extra_args is None else extra_args
723
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
724
+ s_in = x.new_ones([x.shape[0]])
725
+
726
+ for i in trange(len(sigmas) - 1, disable=disable):
727
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
728
+ if callback is not None:
729
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
730
+ x = step_function(x / torch.sqrt(1.0 + sigmas[i] ** 2.0), sigmas[i], sigmas[i + 1], (x - denoised) / sigmas[i], noise_sampler)
731
+ if sigmas[i + 1] != 0:
732
+ x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2.0)
733
+ return x
734
+
735
+
736
+ @torch.no_grad()
737
+ def sample_ddpm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
738
+ return generic_step_sampler(model, x, sigmas, extra_args, callback, disable, noise_sampler, DDPMSampler_step)
739
+
comfy/k_diffusion/utils.py ADDED
@@ -0,0 +1,313 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
14
+
15
+ def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'):
16
+ """Apply passed in transforms for HuggingFace Datasets."""
17
+ images = [transform(image.convert(mode)) for image in examples[image_key]]
18
+ return {image_key: images}
19
+
20
+
21
+ def append_dims(x, target_dims):
22
+ """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
23
+ dims_to_append = target_dims - x.ndim
24
+ if dims_to_append < 0:
25
+ raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
26
+ expanded = x[(...,) + (None,) * dims_to_append]
27
+ # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
28
+ # https://github.com/pytorch/pytorch/issues/84364
29
+ return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
30
+
31
+
32
+ def n_params(module):
33
+ """Returns the number of trainable parameters in a module."""
34
+ return sum(p.numel() for p in module.parameters())
35
+
36
+
37
+ def download_file(path, url, digest=None):
38
+ """Downloads a file if it does not exist, optionally checking its SHA-256 hash."""
39
+ path = Path(path)
40
+ path.parent.mkdir(parents=True, exist_ok=True)
41
+ if not path.exists():
42
+ with urllib.request.urlopen(url) as response, open(path, 'wb') as f:
43
+ shutil.copyfileobj(response, f)
44
+ if digest is not None:
45
+ file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()
46
+ if digest != file_digest:
47
+ raise OSError(f'hash of {path} (url: {url}) failed to validate')
48
+ return path
49
+
50
+
51
+ @contextmanager
52
+ def train_mode(model, mode=True):
53
+ """A context manager that places a model into training mode and restores
54
+ the previous mode on exit."""
55
+ modes = [module.training for module in model.modules()]
56
+ try:
57
+ yield model.train(mode)
58
+ finally:
59
+ for i, module in enumerate(model.modules()):
60
+ module.training = modes[i]
61
+
62
+
63
+ def eval_mode(model):
64
+ """A context manager that places a model into evaluation mode and restores
65
+ the previous mode on exit."""
66
+ return train_mode(model, False)
67
+
68
+
69
+ @torch.no_grad()
70
+ def ema_update(model, averaged_model, decay):
71
+ """Incorporates updated model parameters into an exponential moving averaged
72
+ version of a model. It should be called after each optimizer step."""
73
+ model_params = dict(model.named_parameters())
74
+ averaged_params = dict(averaged_model.named_parameters())
75
+ assert model_params.keys() == averaged_params.keys()
76
+
77
+ for name, param in model_params.items():
78
+ averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)
79
+
80
+ model_buffers = dict(model.named_buffers())
81
+ averaged_buffers = dict(averaged_model.named_buffers())
82
+ assert model_buffers.keys() == averaged_buffers.keys()
83
+
84
+ for name, buf in model_buffers.items():
85
+ averaged_buffers[name].copy_(buf)
86
+
87
+
88
+ class EMAWarmup:
89
+ """Implements an EMA warmup using an inverse decay schedule.
90
+ If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
91
+ good values for models you plan to train for a million or more steps (reaches decay
92
+ factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
93
+ you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
94
+ 215.4k steps).
95
+ Args:
96
+ inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
97
+ power (float): Exponential factor of EMA warmup. Default: 1.
98
+ min_value (float): The minimum EMA decay rate. Default: 0.
99
+ max_value (float): The maximum EMA decay rate. Default: 1.
100
+ start_at (int): The epoch to start averaging at. Default: 0.
101
+ last_epoch (int): The index of last epoch. Default: 0.
102
+ """
103
+
104
+ def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
105
+ last_epoch=0):
106
+ self.inv_gamma = inv_gamma
107
+ self.power = power
108
+ self.min_value = min_value
109
+ self.max_value = max_value
110
+ self.start_at = start_at
111
+ self.last_epoch = last_epoch
112
+
113
+ def state_dict(self):
114
+ """Returns the state of the class as a :class:`dict`."""
115
+ return dict(self.__dict__.items())
116
+
117
+ def load_state_dict(self, state_dict):
118
+ """Loads the class's state.
119
+ Args:
120
+ state_dict (dict): scaler state. Should be an object returned
121
+ from a call to :meth:`state_dict`.
122
+ """
123
+ self.__dict__.update(state_dict)
124
+
125
+ def get_value(self):
126
+ """Gets the current EMA decay rate."""
127
+ epoch = max(0, self.last_epoch - self.start_at)
128
+ value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
129
+ return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
130
+
131
+ def step(self):
132
+ """Updates the step count."""
133
+ self.last_epoch += 1
134
+
135
+
136
+ class InverseLR(optim.lr_scheduler._LRScheduler):
137
+ """Implements an inverse decay learning rate schedule with an optional exponential
138
+ warmup. When last_epoch=-1, sets initial lr as lr.
139
+ inv_gamma is the number of steps/epochs required for the learning rate to decay to
140
+ (1 / 2)**power of its original value.
141
+ Args:
142
+ optimizer (Optimizer): Wrapped optimizer.
143
+ inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
144
+ power (float): Exponential factor of learning rate decay. Default: 1.
145
+ warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
146
+ Default: 0.
147
+ min_lr (float): The minimum learning rate. Default: 0.
148
+ last_epoch (int): The index of last epoch. Default: -1.
149
+ verbose (bool): If ``True``, prints a message to stdout for
150
+ each update. Default: ``False``.
151
+ """
152
+
153
+ def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,
154
+ last_epoch=-1, verbose=False):
155
+ self.inv_gamma = inv_gamma
156
+ self.power = power
157
+ if not 0. <= warmup < 1:
158
+ raise ValueError('Invalid value for warmup')
159
+ self.warmup = warmup
160
+ self.min_lr = min_lr
161
+ super().__init__(optimizer, last_epoch, verbose)
162
+
163
+ def get_lr(self):
164
+ if not self._get_lr_called_within_step:
165
+ warnings.warn("To get the last learning rate computed by the scheduler, "
166
+ "please use `get_last_lr()`.")
167
+
168
+ return self._get_closed_form_lr()
169
+
170
+ def _get_closed_form_lr(self):
171
+ warmup = 1 - self.warmup ** (self.last_epoch + 1)
172
+ lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
173
+ return [warmup * max(self.min_lr, base_lr * lr_mult)
174
+ for base_lr in self.base_lrs]
175
+
176
+
177
+ class ExponentialLR(optim.lr_scheduler._LRScheduler):
178
+ """Implements an exponential learning rate schedule with an optional exponential
179
+ warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate
180
+ continuously by decay (default 0.5) every num_steps steps.
181
+ Args:
182
+ optimizer (Optimizer): Wrapped optimizer.
183
+ num_steps (float): The number of steps to decay the learning rate by decay in.
184
+ decay (float): The factor by which to decay the learning rate every num_steps
185
+ steps. Default: 0.5.
186
+ warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
187
+ Default: 0.
188
+ min_lr (float): The minimum learning rate. Default: 0.
189
+ last_epoch (int): The index of last epoch. Default: -1.
190
+ verbose (bool): If ``True``, prints a message to stdout for
191
+ each update. Default: ``False``.
192
+ """
193
+
194
+ def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,
195
+ last_epoch=-1, verbose=False):
196
+ self.num_steps = num_steps
197
+ self.decay = decay
198
+ if not 0. <= warmup < 1:
199
+ raise ValueError('Invalid value for warmup')
200
+ self.warmup = warmup
201
+ self.min_lr = min_lr
202
+ super().__init__(optimizer, last_epoch, verbose)
203
+
204
+ def get_lr(self):
205
+ if not self._get_lr_called_within_step:
206
+ warnings.warn("To get the last learning rate computed by the scheduler, "
207
+ "please use `get_last_lr()`.")
208
+
209
+ return self._get_closed_form_lr()
210
+
211
+ def _get_closed_form_lr(self):
212
+ warmup = 1 - self.warmup ** (self.last_epoch + 1)
213
+ lr_mult = (self.decay ** (1 / self.num_steps)) ** self.last_epoch
214
+ return [warmup * max(self.min_lr, base_lr * lr_mult)
215
+ for base_lr in self.base_lrs]
216
+
217
+
218
+ def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
219
+ """Draws samples from an lognormal distribution."""
220
+ return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()
221
+
222
+
223
+ def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
224
+ """Draws samples from an optionally truncated log-logistic distribution."""
225
+ min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)
226
+ max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)
227
+ min_cdf = min_value.log().sub(loc).div(scale).sigmoid()
228
+ max_cdf = max_value.log().sub(loc).div(scale).sigmoid()
229
+ u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf
230
+ return u.logit().mul(scale).add(loc).exp().to(dtype)
231
+
232
+
233
+ def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):
234
+ """Draws samples from an log-uniform distribution."""
235
+ min_value = math.log(min_value)
236
+ max_value = math.log(max_value)
237
+ return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()
238
+
239
+
240
+ def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
241
+ """Draws samples from a truncated v-diffusion training timestep distribution."""
242
+ min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi
243
+ max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi
244
+ u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf
245
+ return torch.tan(u * math.pi / 2) * sigma_data
246
+
247
+
248
+ def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32):
249
+ """Draws samples from a split lognormal distribution."""
250
+ n = torch.randn(shape, device=device, dtype=dtype).abs()
251
+ u = torch.rand(shape, device=device, dtype=dtype)
252
+ n_left = n * -scale_1 + loc
253
+ n_right = n * scale_2 + loc
254
+ ratio = scale_1 / (scale_1 + scale_2)
255
+ return torch.where(u < ratio, n_left, n_right).exp()
256
+
257
+
258
+ class FolderOfImages(data.Dataset):
259
+ """Recursively finds all images in a directory. It does not support
260
+ classes/targets."""
261
+
262
+ IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'}
263
+
264
+ def __init__(self, root, transform=None):
265
+ super().__init__()
266
+ self.root = Path(root)
267
+ self.transform = nn.Identity() if transform is None else transform
268
+ self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS)
269
+
270
+ def __repr__(self):
271
+ return f'FolderOfImages(root="{self.root}", len: {len(self)})'
272
+
273
+ def __len__(self):
274
+ return len(self.paths)
275
+
276
+ def __getitem__(self, key):
277
+ path = self.paths[key]
278
+ with open(path, 'rb') as f:
279
+ image = Image.open(f).convert('RGB')
280
+ image = self.transform(image)
281
+ return image,
282
+
283
+
284
+ class CSVLogger:
285
+ def __init__(self, filename, columns):
286
+ self.filename = Path(filename)
287
+ self.columns = columns
288
+ if self.filename.exists():
289
+ self.file = open(self.filename, 'a')
290
+ else:
291
+ self.file = open(self.filename, 'w')
292
+ self.write(*self.columns)
293
+
294
+ def write(self, *args):
295
+ print(*args, sep=',', file=self.file, flush=True)
296
+
297
+
298
+ @contextmanager
299
+ def tf32_mode(cudnn=None, matmul=None):
300
+ """A context manager that sets whether TF32 is allowed on cuDNN or matmul."""
301
+ cudnn_old = torch.backends.cudnn.allow_tf32
302
+ matmul_old = torch.backends.cuda.matmul.allow_tf32
303
+ try:
304
+ if cudnn is not None:
305
+ torch.backends.cudnn.allow_tf32 = cudnn
306
+ if matmul is not None:
307
+ torch.backends.cuda.matmul.allow_tf32 = matmul
308
+ yield
309
+ finally:
310
+ if cudnn is not None:
311
+ torch.backends.cudnn.allow_tf32 = cudnn_old
312
+ if matmul is not None:
313
+ torch.backends.cuda.matmul.allow_tf32 = matmul_old
comfy/latent_formats.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ class LatentFormat:
3
+ scale_factor = 1.0
4
+ latent_rgb_factors = None
5
+ taesd_decoder_name = None
6
+
7
+ def process_in(self, latent):
8
+ return latent * self.scale_factor
9
+
10
+ def process_out(self, latent):
11
+ return latent / self.scale_factor
12
+
13
+ class SD15(LatentFormat):
14
+ def __init__(self, scale_factor=0.18215):
15
+ self.scale_factor = scale_factor
16
+ self.latent_rgb_factors = [
17
+ # R G B
18
+ [ 0.3512, 0.2297, 0.3227],
19
+ [ 0.3250, 0.4974, 0.2350],
20
+ [-0.2829, 0.1762, 0.2721],
21
+ [-0.2120, -0.2616, -0.7177]
22
+ ]
23
+ self.taesd_decoder_name = "taesd_decoder.pth"
24
+
25
+ class SDXL(LatentFormat):
26
+ def __init__(self):
27
+ self.scale_factor = 0.13025
28
+ self.latent_rgb_factors = [
29
+ # R G B
30
+ [ 0.3920, 0.4054, 0.4549],
31
+ [-0.2634, -0.0196, 0.0653],
32
+ [ 0.0568, 0.1687, -0.0755],
33
+ [-0.3112, -0.2359, -0.2076]
34
+ ]
35
+ self.taesd_decoder_name = "taesdxl_decoder.pth"
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 comfy.ldm.modules.diffusionmodules.model import Encoder, Decoder
7
+ from comfy.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
8
+
9
+ from comfy.ldm.util import instantiate_from_config
10
+ from comfy.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,418 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SAMPLING ONLY."""
2
+
3
+ import torch
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+
7
+ from comfy.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
+ self.parameterization = kwargs.get("parameterization", "eps")
18
+
19
+ def register_buffer(self, name, attr):
20
+ if type(attr) == torch.Tensor:
21
+ if attr.device != self.device:
22
+ attr = attr.float().to(self.device)
23
+ setattr(self, name, attr)
24
+
25
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
26
+ ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
27
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
28
+ self.make_schedule_timesteps(ddim_timesteps, ddim_eta=ddim_eta, verbose=verbose)
29
+
30
+ def make_schedule_timesteps(self, ddim_timesteps, ddim_eta=0., verbose=True):
31
+ self.ddim_timesteps = torch.tensor(ddim_timesteps)
32
+ alphas_cumprod = self.model.alphas_cumprod
33
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
34
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.device)
35
+
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
+ disable_pbar=False,
85
+ **kwargs
86
+ ):
87
+ self.make_schedule_timesteps(ddim_timesteps=ddim_timesteps, ddim_eta=eta, verbose=verbose)
88
+ samples, intermediates = self.ddim_sampling(conditioning, x_T.shape,
89
+ callback=callback,
90
+ img_callback=img_callback,
91
+ quantize_denoised=quantize_x0,
92
+ mask=mask, x0=x0,
93
+ ddim_use_original_steps=False,
94
+ noise_dropout=noise_dropout,
95
+ temperature=temperature,
96
+ score_corrector=score_corrector,
97
+ corrector_kwargs=corrector_kwargs,
98
+ x_T=x_T,
99
+ log_every_t=log_every_t,
100
+ unconditional_guidance_scale=unconditional_guidance_scale,
101
+ unconditional_conditioning=unconditional_conditioning,
102
+ dynamic_threshold=dynamic_threshold,
103
+ ucg_schedule=ucg_schedule,
104
+ denoise_function=denoise_function,
105
+ extra_args=extra_args,
106
+ to_zero=to_zero,
107
+ end_step=end_step,
108
+ disable_pbar=disable_pbar
109
+ )
110
+ return samples, intermediates
111
+
112
+
113
+ @torch.no_grad()
114
+ def sample(self,
115
+ S,
116
+ batch_size,
117
+ shape,
118
+ conditioning=None,
119
+ callback=None,
120
+ normals_sequence=None,
121
+ img_callback=None,
122
+ quantize_x0=False,
123
+ eta=0.,
124
+ mask=None,
125
+ x0=None,
126
+ temperature=1.,
127
+ noise_dropout=0.,
128
+ score_corrector=None,
129
+ corrector_kwargs=None,
130
+ verbose=True,
131
+ x_T=None,
132
+ log_every_t=100,
133
+ unconditional_guidance_scale=1.,
134
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
135
+ dynamic_threshold=None,
136
+ ucg_schedule=None,
137
+ **kwargs
138
+ ):
139
+ if conditioning is not None:
140
+ if isinstance(conditioning, dict):
141
+ ctmp = conditioning[list(conditioning.keys())[0]]
142
+ while isinstance(ctmp, list): ctmp = ctmp[0]
143
+ cbs = ctmp.shape[0]
144
+ if cbs != batch_size:
145
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
146
+
147
+ elif isinstance(conditioning, list):
148
+ for ctmp in conditioning:
149
+ if ctmp.shape[0] != batch_size:
150
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
151
+
152
+ else:
153
+ if conditioning.shape[0] != batch_size:
154
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
155
+
156
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
157
+ # sampling
158
+ C, H, W = shape
159
+ size = (batch_size, C, H, W)
160
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
161
+
162
+ samples, intermediates = self.ddim_sampling(conditioning, size,
163
+ callback=callback,
164
+ img_callback=img_callback,
165
+ quantize_denoised=quantize_x0,
166
+ mask=mask, x0=x0,
167
+ ddim_use_original_steps=False,
168
+ noise_dropout=noise_dropout,
169
+ temperature=temperature,
170
+ score_corrector=score_corrector,
171
+ corrector_kwargs=corrector_kwargs,
172
+ x_T=x_T,
173
+ log_every_t=log_every_t,
174
+ unconditional_guidance_scale=unconditional_guidance_scale,
175
+ unconditional_conditioning=unconditional_conditioning,
176
+ dynamic_threshold=dynamic_threshold,
177
+ ucg_schedule=ucg_schedule,
178
+ denoise_function=None,
179
+ extra_args=None
180
+ )
181
+ return samples, intermediates
182
+
183
+ def q_sample(self, x_start, t, noise=None):
184
+ if noise is None:
185
+ noise = torch.randn_like(x_start)
186
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
187
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
188
+
189
+ @torch.no_grad()
190
+ def ddim_sampling(self, cond, shape,
191
+ x_T=None, ddim_use_original_steps=False,
192
+ callback=None, timesteps=None, quantize_denoised=False,
193
+ mask=None, x0=None, img_callback=None, log_every_t=100,
194
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
195
+ unconditional_guidance_scale=1., unconditional_conditioning=None, dynamic_threshold=None,
196
+ ucg_schedule=None, denoise_function=None, extra_args=None, to_zero=True, end_step=None, disable_pbar=False):
197
+ device = self.model.alphas_cumprod.device
198
+ b = shape[0]
199
+ if x_T is None:
200
+ img = torch.randn(shape, device=device)
201
+ else:
202
+ img = x_T
203
+
204
+ if timesteps is None:
205
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
206
+ elif timesteps is not None and not ddim_use_original_steps:
207
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
208
+ timesteps = self.ddim_timesteps[:subset_end]
209
+
210
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
211
+ time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else timesteps.flip(0)
212
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
213
+ # print(f"Running DDIM Sampling with {total_steps} timesteps")
214
+
215
+ iterator = tqdm(time_range[:end_step], desc='DDIM Sampler', total=end_step, disable=disable_pbar)
216
+
217
+ for i, step in enumerate(iterator):
218
+ index = total_steps - i - 1
219
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
220
+
221
+ if mask is not None:
222
+ assert x0 is not None
223
+ img_orig = self.q_sample(x0, ts) # TODO: deterministic forward pass?
224
+ img = img_orig * mask + (1. - mask) * img
225
+
226
+ if ucg_schedule is not None:
227
+ assert len(ucg_schedule) == len(time_range)
228
+ unconditional_guidance_scale = ucg_schedule[i]
229
+
230
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
231
+ quantize_denoised=quantize_denoised, temperature=temperature,
232
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
233
+ corrector_kwargs=corrector_kwargs,
234
+ unconditional_guidance_scale=unconditional_guidance_scale,
235
+ unconditional_conditioning=unconditional_conditioning,
236
+ dynamic_threshold=dynamic_threshold, denoise_function=denoise_function, extra_args=extra_args)
237
+ img, pred_x0 = outs
238
+ if callback: callback(i)
239
+ if img_callback: img_callback(pred_x0, i)
240
+
241
+ if index % log_every_t == 0 or index == total_steps - 1:
242
+ intermediates['x_inter'].append(img)
243
+ intermediates['pred_x0'].append(pred_x0)
244
+
245
+ if to_zero:
246
+ img = pred_x0
247
+ else:
248
+ if ddim_use_original_steps:
249
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
250
+ else:
251
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
252
+ img /= sqrt_alphas_cumprod[index - 1]
253
+
254
+ return img, intermediates
255
+
256
+ @torch.no_grad()
257
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
258
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
259
+ unconditional_guidance_scale=1., unconditional_conditioning=None,
260
+ dynamic_threshold=None, denoise_function=None, extra_args=None):
261
+ b, *_, device = *x.shape, x.device
262
+
263
+ if denoise_function is not None:
264
+ model_output = denoise_function(x, t, **extra_args)
265
+ elif unconditional_conditioning is None or unconditional_guidance_scale == 1.:
266
+ model_output = self.model.apply_model(x, t, c)
267
+ else:
268
+ x_in = torch.cat([x] * 2)
269
+ t_in = torch.cat([t] * 2)
270
+ if isinstance(c, dict):
271
+ assert isinstance(unconditional_conditioning, dict)
272
+ c_in = dict()
273
+ for k in c:
274
+ if isinstance(c[k], list):
275
+ c_in[k] = [torch.cat([
276
+ unconditional_conditioning[k][i],
277
+ c[k][i]]) for i in range(len(c[k]))]
278
+ else:
279
+ c_in[k] = torch.cat([
280
+ unconditional_conditioning[k],
281
+ c[k]])
282
+ elif isinstance(c, list):
283
+ c_in = list()
284
+ assert isinstance(unconditional_conditioning, list)
285
+ for i in range(len(c)):
286
+ c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
287
+ else:
288
+ c_in = torch.cat([unconditional_conditioning, c])
289
+ model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
290
+ model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
291
+
292
+ if self.parameterization == "v":
293
+ e_t = extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * model_output + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
294
+ else:
295
+ e_t = model_output
296
+
297
+ if score_corrector is not None:
298
+ assert self.parameterization == "eps", 'not implemented'
299
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
300
+
301
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
302
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
303
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
304
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
305
+ # select parameters corresponding to the currently considered timestep
306
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
307
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
308
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
309
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
310
+
311
+ # current prediction for x_0
312
+ if self.parameterization != "v":
313
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
314
+ else:
315
+ pred_x0 = extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * x - extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * model_output
316
+
317
+ if quantize_denoised:
318
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
319
+
320
+ if dynamic_threshold is not None:
321
+ raise NotImplementedError()
322
+
323
+ # direction pointing to x_t
324
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
325
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
326
+ if noise_dropout > 0.:
327
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
328
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
329
+ return x_prev, pred_x0
330
+
331
+ @torch.no_grad()
332
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
333
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None):
334
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
335
+
336
+ assert t_enc <= num_reference_steps
337
+ num_steps = t_enc
338
+
339
+ if use_original_steps:
340
+ alphas_next = self.alphas_cumprod[:num_steps]
341
+ alphas = self.alphas_cumprod_prev[:num_steps]
342
+ else:
343
+ alphas_next = self.ddim_alphas[:num_steps]
344
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
345
+
346
+ x_next = x0
347
+ intermediates = []
348
+ inter_steps = []
349
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
350
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
351
+ if unconditional_guidance_scale == 1.:
352
+ noise_pred = self.model.apply_model(x_next, t, c)
353
+ else:
354
+ assert unconditional_conditioning is not None
355
+ e_t_uncond, noise_pred = torch.chunk(
356
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
357
+ torch.cat((unconditional_conditioning, c))), 2)
358
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
359
+
360
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
361
+ weighted_noise_pred = alphas_next[i].sqrt() * (
362
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
363
+ x_next = xt_weighted + weighted_noise_pred
364
+ if return_intermediates and i % (
365
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
366
+ intermediates.append(x_next)
367
+ inter_steps.append(i)
368
+ elif return_intermediates and i >= num_steps - 2:
369
+ intermediates.append(x_next)
370
+ inter_steps.append(i)
371
+ if callback: callback(i)
372
+
373
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
374
+ if return_intermediates:
375
+ out.update({'intermediates': intermediates})
376
+ return x_next, out
377
+
378
+ @torch.no_grad()
379
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None, max_denoise=False):
380
+ # fast, but does not allow for exact reconstruction
381
+ # t serves as an index to gather the correct alphas
382
+ if use_original_steps:
383
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
384
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
385
+ else:
386
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
387
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
388
+
389
+ if noise is None:
390
+ noise = torch.randn_like(x0)
391
+ if max_denoise:
392
+ noise_multiplier = 1.0
393
+ else:
394
+ noise_multiplier = extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape)
395
+
396
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 + noise_multiplier * noise)
397
+
398
+ @torch.no_grad()
399
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
400
+ use_original_steps=False, callback=None):
401
+
402
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
403
+ timesteps = timesteps[:t_start]
404
+
405
+ time_range = np.flip(timesteps)
406
+ total_steps = timesteps.shape[0]
407
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
408
+
409
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
410
+ x_dec = x_latent
411
+ for i, step in enumerate(iterator):
412
+ index = total_steps - i - 1
413
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
414
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
415
+ unconditional_guidance_scale=unconditional_guidance_scale,
416
+ unconditional_conditioning=unconditional_conditioning)
417
+ if callback: callback(i)
418
+ return x_dec
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,700 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 .diffusionmodules.util import checkpoint
10
+ from .sub_quadratic_attention import efficient_dot_product_attention
11
+
12
+ from comfy import model_management
13
+
14
+ if model_management.xformers_enabled():
15
+ import xformers
16
+ import xformers.ops
17
+
18
+ from comfy.cli_args import args
19
+ import comfy.ops
20
+
21
+ # CrossAttn precision handling
22
+ if args.dont_upcast_attention:
23
+ print("disabling upcasting of attention")
24
+ _ATTN_PRECISION = "fp16"
25
+ else:
26
+ _ATTN_PRECISION = "fp32"
27
+
28
+
29
+ def exists(val):
30
+ return val is not None
31
+
32
+
33
+ def uniq(arr):
34
+ return{el: True for el in arr}.keys()
35
+
36
+
37
+ def default(val, d):
38
+ if exists(val):
39
+ return val
40
+ return d
41
+
42
+
43
+ def max_neg_value(t):
44
+ return -torch.finfo(t.dtype).max
45
+
46
+
47
+ def init_(tensor):
48
+ dim = tensor.shape[-1]
49
+ std = 1 / math.sqrt(dim)
50
+ tensor.uniform_(-std, std)
51
+ return tensor
52
+
53
+
54
+ # feedforward
55
+ class GEGLU(nn.Module):
56
+ def __init__(self, dim_in, dim_out, dtype=None, device=None, operations=comfy.ops):
57
+ super().__init__()
58
+ self.proj = operations.Linear(dim_in, dim_out * 2, dtype=dtype, device=device)
59
+
60
+ def forward(self, x):
61
+ x, gate = self.proj(x).chunk(2, dim=-1)
62
+ return x * F.gelu(gate)
63
+
64
+
65
+ class FeedForward(nn.Module):
66
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., dtype=None, device=None, operations=comfy.ops):
67
+ super().__init__()
68
+ inner_dim = int(dim * mult)
69
+ dim_out = default(dim_out, dim)
70
+ project_in = nn.Sequential(
71
+ operations.Linear(dim, inner_dim, dtype=dtype, device=device),
72
+ nn.GELU()
73
+ ) if not glu else GEGLU(dim, inner_dim, dtype=dtype, device=device, operations=operations)
74
+
75
+ self.net = nn.Sequential(
76
+ project_in,
77
+ nn.Dropout(dropout),
78
+ operations.Linear(inner_dim, dim_out, dtype=dtype, device=device)
79
+ )
80
+
81
+ def forward(self, x):
82
+ return self.net(x)
83
+
84
+
85
+ def zero_module(module):
86
+ """
87
+ Zero out the parameters of a module and return it.
88
+ """
89
+ for p in module.parameters():
90
+ p.detach().zero_()
91
+ return module
92
+
93
+
94
+ def Normalize(in_channels, dtype=None, device=None):
95
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
96
+
97
+
98
+ class SpatialSelfAttention(nn.Module):
99
+ def __init__(self, in_channels):
100
+ super().__init__()
101
+ self.in_channels = in_channels
102
+
103
+ self.norm = Normalize(in_channels)
104
+ self.q = torch.nn.Conv2d(in_channels,
105
+ in_channels,
106
+ kernel_size=1,
107
+ stride=1,
108
+ padding=0)
109
+ self.k = torch.nn.Conv2d(in_channels,
110
+ in_channels,
111
+ kernel_size=1,
112
+ stride=1,
113
+ padding=0)
114
+ self.v = torch.nn.Conv2d(in_channels,
115
+ in_channels,
116
+ kernel_size=1,
117
+ stride=1,
118
+ padding=0)
119
+ self.proj_out = torch.nn.Conv2d(in_channels,
120
+ in_channels,
121
+ kernel_size=1,
122
+ stride=1,
123
+ padding=0)
124
+
125
+ def forward(self, x):
126
+ h_ = x
127
+ h_ = self.norm(h_)
128
+ q = self.q(h_)
129
+ k = self.k(h_)
130
+ v = self.v(h_)
131
+
132
+ # compute attention
133
+ b,c,h,w = q.shape
134
+ q = rearrange(q, 'b c h w -> b (h w) c')
135
+ k = rearrange(k, 'b c h w -> b c (h w)')
136
+ w_ = torch.einsum('bij,bjk->bik', q, k)
137
+
138
+ w_ = w_ * (int(c)**(-0.5))
139
+ w_ = torch.nn.functional.softmax(w_, dim=2)
140
+
141
+ # attend to values
142
+ v = rearrange(v, 'b c h w -> b c (h w)')
143
+ w_ = rearrange(w_, 'b i j -> b j i')
144
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
145
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
146
+ h_ = self.proj_out(h_)
147
+
148
+ return x+h_
149
+
150
+
151
+ class CrossAttentionBirchSan(nn.Module):
152
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops):
153
+ super().__init__()
154
+ inner_dim = dim_head * heads
155
+ context_dim = default(context_dim, query_dim)
156
+
157
+ self.scale = dim_head ** -0.5
158
+ self.heads = heads
159
+
160
+ self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
161
+ self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
162
+ self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
163
+
164
+ self.to_out = nn.Sequential(
165
+ operations.Linear(inner_dim, query_dim, dtype=dtype, device=device),
166
+ nn.Dropout(dropout)
167
+ )
168
+
169
+ def forward(self, x, context=None, value=None, mask=None):
170
+ h = self.heads
171
+
172
+ query = self.to_q(x)
173
+ context = default(context, x)
174
+ key = self.to_k(context)
175
+ if value is not None:
176
+ value = self.to_v(value)
177
+ else:
178
+ value = self.to_v(context)
179
+
180
+ del context, x
181
+
182
+ query = query.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1)
183
+ key_t = key.transpose(1,2).unflatten(1, (self.heads, -1)).flatten(end_dim=1)
184
+ del key
185
+ value = value.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1)
186
+
187
+ dtype = query.dtype
188
+ upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32
189
+ if upcast_attention:
190
+ bytes_per_token = torch.finfo(torch.float32).bits//8
191
+ else:
192
+ bytes_per_token = torch.finfo(query.dtype).bits//8
193
+ batch_x_heads, q_tokens, _ = query.shape
194
+ _, _, k_tokens = key_t.shape
195
+ qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
196
+
197
+ mem_free_total, mem_free_torch = model_management.get_free_memory(query.device, True)
198
+
199
+ chunk_threshold_bytes = mem_free_torch * 0.5 #Using only this seems to work better on AMD
200
+
201
+ kv_chunk_size_min = None
202
+
203
+ #not sure at all about the math here
204
+ #TODO: tweak this
205
+ if mem_free_total > 8192 * 1024 * 1024 * 1.3:
206
+ query_chunk_size_x = 1024 * 4
207
+ elif mem_free_total > 4096 * 1024 * 1024 * 1.3:
208
+ query_chunk_size_x = 1024 * 2
209
+ else:
210
+ query_chunk_size_x = 1024
211
+ kv_chunk_size_min_x = None
212
+ kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024
213
+ if kv_chunk_size_x < 1024:
214
+ kv_chunk_size_x = None
215
+
216
+ if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes:
217
+ # the big matmul fits into our memory limit; do everything in 1 chunk,
218
+ # i.e. send it down the unchunked fast-path
219
+ query_chunk_size = q_tokens
220
+ kv_chunk_size = k_tokens
221
+ else:
222
+ query_chunk_size = query_chunk_size_x
223
+ kv_chunk_size = kv_chunk_size_x
224
+ kv_chunk_size_min = kv_chunk_size_min_x
225
+
226
+ hidden_states = efficient_dot_product_attention(
227
+ query,
228
+ key_t,
229
+ value,
230
+ query_chunk_size=query_chunk_size,
231
+ kv_chunk_size=kv_chunk_size,
232
+ kv_chunk_size_min=kv_chunk_size_min,
233
+ use_checkpoint=self.training,
234
+ upcast_attention=upcast_attention,
235
+ )
236
+
237
+ hidden_states = hidden_states.to(dtype)
238
+
239
+ hidden_states = hidden_states.unflatten(0, (-1, self.heads)).transpose(1,2).flatten(start_dim=2)
240
+
241
+ out_proj, dropout = self.to_out
242
+ hidden_states = out_proj(hidden_states)
243
+ hidden_states = dropout(hidden_states)
244
+
245
+ return hidden_states
246
+
247
+
248
+ class CrossAttentionDoggettx(nn.Module):
249
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops):
250
+ super().__init__()
251
+ inner_dim = dim_head * heads
252
+ context_dim = default(context_dim, query_dim)
253
+
254
+ self.scale = dim_head ** -0.5
255
+ self.heads = heads
256
+
257
+ self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
258
+ self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
259
+ self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
260
+
261
+ self.to_out = nn.Sequential(
262
+ operations.Linear(inner_dim, query_dim, dtype=dtype, device=device),
263
+ nn.Dropout(dropout)
264
+ )
265
+
266
+ def forward(self, x, context=None, value=None, mask=None):
267
+ h = self.heads
268
+
269
+ q_in = self.to_q(x)
270
+ context = default(context, x)
271
+ k_in = self.to_k(context)
272
+ if value is not None:
273
+ v_in = self.to_v(value)
274
+ del value
275
+ else:
276
+ v_in = self.to_v(context)
277
+ del context, x
278
+
279
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
280
+ del q_in, k_in, v_in
281
+
282
+ r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
283
+
284
+ mem_free_total = model_management.get_free_memory(q.device)
285
+
286
+ gb = 1024 ** 3
287
+ tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
288
+ modifier = 3 if q.element_size() == 2 else 2.5
289
+ mem_required = tensor_size * modifier
290
+ steps = 1
291
+
292
+
293
+ if mem_required > mem_free_total:
294
+ steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
295
+ # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
296
+ # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
297
+
298
+ if steps > 64:
299
+ max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
300
+ raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
301
+ f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
302
+
303
+ # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size)
304
+ first_op_done = False
305
+ cleared_cache = False
306
+ while True:
307
+ try:
308
+ slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
309
+ for i in range(0, q.shape[1], slice_size):
310
+ end = i + slice_size
311
+ if _ATTN_PRECISION =="fp32":
312
+ with torch.autocast(enabled=False, device_type = 'cuda'):
313
+ s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * self.scale
314
+ else:
315
+ s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
316
+ first_op_done = True
317
+
318
+ s2 = s1.softmax(dim=-1).to(v.dtype)
319
+ del s1
320
+
321
+ r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
322
+ del s2
323
+ break
324
+ except model_management.OOM_EXCEPTION as e:
325
+ if first_op_done == False:
326
+ model_management.soft_empty_cache(True)
327
+ if cleared_cache == False:
328
+ cleared_cache = True
329
+ print("out of memory error, emptying cache and trying again")
330
+ continue
331
+ steps *= 2
332
+ if steps > 64:
333
+ raise e
334
+ print("out of memory error, increasing steps and trying again", steps)
335
+ else:
336
+ raise e
337
+
338
+ del q, k, v
339
+
340
+ r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
341
+ del r1
342
+
343
+ return self.to_out(r2)
344
+
345
+ class CrossAttention(nn.Module):
346
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops):
347
+ super().__init__()
348
+ inner_dim = dim_head * heads
349
+ context_dim = default(context_dim, query_dim)
350
+
351
+ self.scale = dim_head ** -0.5
352
+ self.heads = heads
353
+
354
+ self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
355
+ self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
356
+ self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
357
+
358
+ self.to_out = nn.Sequential(
359
+ operations.Linear(inner_dim, query_dim, dtype=dtype, device=device),
360
+ nn.Dropout(dropout)
361
+ )
362
+
363
+ def forward(self, x, context=None, value=None, mask=None):
364
+ h = self.heads
365
+
366
+ q = self.to_q(x)
367
+ context = default(context, x)
368
+ k = self.to_k(context)
369
+ if value is not None:
370
+ v = self.to_v(value)
371
+ del value
372
+ else:
373
+ v = self.to_v(context)
374
+
375
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
376
+
377
+ # force cast to fp32 to avoid overflowing
378
+ if _ATTN_PRECISION =="fp32":
379
+ with torch.autocast(enabled=False, device_type = 'cuda'):
380
+ q, k = q.float(), k.float()
381
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
382
+ else:
383
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
384
+
385
+ del q, k
386
+
387
+ if exists(mask):
388
+ mask = rearrange(mask, 'b ... -> b (...)')
389
+ max_neg_value = -torch.finfo(sim.dtype).max
390
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
391
+ sim.masked_fill_(~mask, max_neg_value)
392
+
393
+ # attention, what we cannot get enough of
394
+ sim = sim.softmax(dim=-1)
395
+
396
+ out = einsum('b i j, b j d -> b i d', sim, v)
397
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
398
+ return self.to_out(out)
399
+
400
+ class MemoryEfficientCrossAttention(nn.Module):
401
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
402
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, dtype=None, device=None, operations=comfy.ops):
403
+ super().__init__()
404
+ inner_dim = dim_head * heads
405
+ context_dim = default(context_dim, query_dim)
406
+
407
+ self.heads = heads
408
+ self.dim_head = dim_head
409
+
410
+ self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
411
+ self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
412
+ self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
413
+
414
+ self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
415
+ self.attention_op: Optional[Any] = None
416
+
417
+ def forward(self, x, context=None, value=None, mask=None):
418
+ q = self.to_q(x)
419
+ context = default(context, x)
420
+ k = self.to_k(context)
421
+ if value is not None:
422
+ v = self.to_v(value)
423
+ del value
424
+ else:
425
+ v = self.to_v(context)
426
+
427
+ b, _, _ = q.shape
428
+ q, k, v = map(
429
+ lambda t: t.unsqueeze(3)
430
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
431
+ .permute(0, 2, 1, 3)
432
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
433
+ .contiguous(),
434
+ (q, k, v),
435
+ )
436
+
437
+ # actually compute the attention, what we cannot get enough of
438
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
439
+
440
+ if exists(mask):
441
+ raise NotImplementedError
442
+ out = (
443
+ out.unsqueeze(0)
444
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
445
+ .permute(0, 2, 1, 3)
446
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
447
+ )
448
+ return self.to_out(out)
449
+
450
+ class CrossAttentionPytorch(nn.Module):
451
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops):
452
+ super().__init__()
453
+ inner_dim = dim_head * heads
454
+ context_dim = default(context_dim, query_dim)
455
+
456
+ self.heads = heads
457
+ self.dim_head = dim_head
458
+
459
+ self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
460
+ self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
461
+ self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
462
+
463
+ self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
464
+ self.attention_op: Optional[Any] = None
465
+
466
+ def forward(self, x, context=None, value=None, mask=None):
467
+ q = self.to_q(x)
468
+ context = default(context, x)
469
+ k = self.to_k(context)
470
+ if value is not None:
471
+ v = self.to_v(value)
472
+ del value
473
+ else:
474
+ v = self.to_v(context)
475
+
476
+ b, _, _ = q.shape
477
+ q, k, v = map(
478
+ lambda t: t.view(b, -1, self.heads, self.dim_head).transpose(1, 2),
479
+ (q, k, v),
480
+ )
481
+
482
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
483
+
484
+ if exists(mask):
485
+ raise NotImplementedError
486
+ out = (
487
+ out.transpose(1, 2).reshape(b, -1, self.heads * self.dim_head)
488
+ )
489
+
490
+ return self.to_out(out)
491
+
492
+ if model_management.xformers_enabled():
493
+ print("Using xformers cross attention")
494
+ CrossAttention = MemoryEfficientCrossAttention
495
+ elif model_management.pytorch_attention_enabled():
496
+ print("Using pytorch cross attention")
497
+ CrossAttention = CrossAttentionPytorch
498
+ else:
499
+ if args.use_split_cross_attention:
500
+ print("Using split optimization for cross attention")
501
+ CrossAttention = CrossAttentionDoggettx
502
+ else:
503
+ print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention")
504
+ CrossAttention = CrossAttentionBirchSan
505
+
506
+
507
+ class BasicTransformerBlock(nn.Module):
508
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
509
+ disable_self_attn=False, dtype=None, device=None, operations=comfy.ops):
510
+ super().__init__()
511
+ self.disable_self_attn = disable_self_attn
512
+ self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
513
+ context_dim=context_dim if self.disable_self_attn else None, dtype=dtype, device=device, operations=operations) # is a self-attention if not self.disable_self_attn
514
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff, dtype=dtype, device=device, operations=operations)
515
+ self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
516
+ heads=n_heads, dim_head=d_head, dropout=dropout, dtype=dtype, device=device, operations=operations) # is self-attn if context is none
517
+ self.norm1 = nn.LayerNorm(dim, dtype=dtype, device=device)
518
+ self.norm2 = nn.LayerNorm(dim, dtype=dtype, device=device)
519
+ self.norm3 = nn.LayerNorm(dim, dtype=dtype, device=device)
520
+ self.checkpoint = checkpoint
521
+ self.n_heads = n_heads
522
+ self.d_head = d_head
523
+
524
+ def forward(self, x, context=None, transformer_options={}):
525
+ return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint)
526
+
527
+ def _forward(self, x, context=None, transformer_options={}):
528
+ extra_options = {}
529
+ block = None
530
+ block_index = 0
531
+ if "current_index" in transformer_options:
532
+ extra_options["transformer_index"] = transformer_options["current_index"]
533
+ if "block_index" in transformer_options:
534
+ block_index = transformer_options["block_index"]
535
+ extra_options["block_index"] = block_index
536
+ if "original_shape" in transformer_options:
537
+ extra_options["original_shape"] = transformer_options["original_shape"]
538
+ if "block" in transformer_options:
539
+ block = transformer_options["block"]
540
+ extra_options["block"] = block
541
+ if "patches" in transformer_options:
542
+ transformer_patches = transformer_options["patches"]
543
+ else:
544
+ transformer_patches = {}
545
+
546
+ extra_options["n_heads"] = self.n_heads
547
+ extra_options["dim_head"] = self.d_head
548
+
549
+ if "patches_replace" in transformer_options:
550
+ transformer_patches_replace = transformer_options["patches_replace"]
551
+ else:
552
+ transformer_patches_replace = {}
553
+
554
+ n = self.norm1(x)
555
+ if self.disable_self_attn:
556
+ context_attn1 = context
557
+ else:
558
+ context_attn1 = None
559
+ value_attn1 = None
560
+
561
+ if "attn1_patch" in transformer_patches:
562
+ patch = transformer_patches["attn1_patch"]
563
+ if context_attn1 is None:
564
+ context_attn1 = n
565
+ value_attn1 = context_attn1
566
+ for p in patch:
567
+ n, context_attn1, value_attn1 = p(n, context_attn1, value_attn1, extra_options)
568
+
569
+ if block is not None:
570
+ transformer_block = (block[0], block[1], block_index)
571
+ else:
572
+ transformer_block = None
573
+ attn1_replace_patch = transformer_patches_replace.get("attn1", {})
574
+ block_attn1 = transformer_block
575
+ if block_attn1 not in attn1_replace_patch:
576
+ block_attn1 = block
577
+
578
+ if block_attn1 in attn1_replace_patch:
579
+ if context_attn1 is None:
580
+ context_attn1 = n
581
+ value_attn1 = n
582
+ n = self.attn1.to_q(n)
583
+ context_attn1 = self.attn1.to_k(context_attn1)
584
+ value_attn1 = self.attn1.to_v(value_attn1)
585
+ n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options)
586
+ n = self.attn1.to_out(n)
587
+ else:
588
+ n = self.attn1(n, context=context_attn1, value=value_attn1)
589
+
590
+ if "attn1_output_patch" in transformer_patches:
591
+ patch = transformer_patches["attn1_output_patch"]
592
+ for p in patch:
593
+ n = p(n, extra_options)
594
+
595
+ x += n
596
+ if "middle_patch" in transformer_patches:
597
+ patch = transformer_patches["middle_patch"]
598
+ for p in patch:
599
+ x = p(x, extra_options)
600
+
601
+ n = self.norm2(x)
602
+
603
+ context_attn2 = context
604
+ value_attn2 = None
605
+ if "attn2_patch" in transformer_patches:
606
+ patch = transformer_patches["attn2_patch"]
607
+ value_attn2 = context_attn2
608
+ for p in patch:
609
+ n, context_attn2, value_attn2 = p(n, context_attn2, value_attn2, extra_options)
610
+
611
+ attn2_replace_patch = transformer_patches_replace.get("attn2", {})
612
+ block_attn2 = transformer_block
613
+ if block_attn2 not in attn2_replace_patch:
614
+ block_attn2 = block
615
+
616
+ if block_attn2 in attn2_replace_patch:
617
+ if value_attn2 is None:
618
+ value_attn2 = context_attn2
619
+ n = self.attn2.to_q(n)
620
+ context_attn2 = self.attn2.to_k(context_attn2)
621
+ value_attn2 = self.attn2.to_v(value_attn2)
622
+ n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options)
623
+ n = self.attn2.to_out(n)
624
+ else:
625
+ n = self.attn2(n, context=context_attn2, value=value_attn2)
626
+
627
+ if "attn2_output_patch" in transformer_patches:
628
+ patch = transformer_patches["attn2_output_patch"]
629
+ for p in patch:
630
+ n = p(n, extra_options)
631
+
632
+ x += n
633
+ x = self.ff(self.norm3(x)) + x
634
+ return x
635
+
636
+
637
+ class SpatialTransformer(nn.Module):
638
+ """
639
+ Transformer block for image-like data.
640
+ First, project the input (aka embedding)
641
+ and reshape to b, t, d.
642
+ Then apply standard transformer action.
643
+ Finally, reshape to image
644
+ NEW: use_linear for more efficiency instead of the 1x1 convs
645
+ """
646
+ def __init__(self, in_channels, n_heads, d_head,
647
+ depth=1, dropout=0., context_dim=None,
648
+ disable_self_attn=False, use_linear=False,
649
+ use_checkpoint=True, dtype=None, device=None, operations=comfy.ops):
650
+ super().__init__()
651
+ if exists(context_dim) and not isinstance(context_dim, list):
652
+ context_dim = [context_dim] * depth
653
+ self.in_channels = in_channels
654
+ inner_dim = n_heads * d_head
655
+ self.norm = Normalize(in_channels, dtype=dtype, device=device)
656
+ if not use_linear:
657
+ self.proj_in = operations.Conv2d(in_channels,
658
+ inner_dim,
659
+ kernel_size=1,
660
+ stride=1,
661
+ padding=0, dtype=dtype, device=device)
662
+ else:
663
+ self.proj_in = operations.Linear(in_channels, inner_dim, dtype=dtype, device=device)
664
+
665
+ self.transformer_blocks = nn.ModuleList(
666
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
667
+ disable_self_attn=disable_self_attn, checkpoint=use_checkpoint, dtype=dtype, device=device, operations=operations)
668
+ for d in range(depth)]
669
+ )
670
+ if not use_linear:
671
+ self.proj_out = operations.Conv2d(inner_dim,in_channels,
672
+ kernel_size=1,
673
+ stride=1,
674
+ padding=0, dtype=dtype, device=device)
675
+ else:
676
+ self.proj_out = operations.Linear(in_channels, inner_dim, dtype=dtype, device=device)
677
+ self.use_linear = use_linear
678
+
679
+ def forward(self, x, context=None, transformer_options={}):
680
+ # note: if no context is given, cross-attention defaults to self-attention
681
+ if not isinstance(context, list):
682
+ context = [context] * len(self.transformer_blocks)
683
+ b, c, h, w = x.shape
684
+ x_in = x
685
+ x = self.norm(x)
686
+ if not self.use_linear:
687
+ x = self.proj_in(x)
688
+ x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
689
+ if self.use_linear:
690
+ x = self.proj_in(x)
691
+ for i, block in enumerate(self.transformer_blocks):
692
+ transformer_options["block_index"] = i
693
+ x = block(x, context=context[i], transformer_options=transformer_options)
694
+ if self.use_linear:
695
+ x = self.proj_out(x)
696
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
697
+ if not self.use_linear:
698
+ x = self.proj_out(x)
699
+ return x + x_in
700
+
comfy/ldm/modules/diffusionmodules/__init__.py ADDED
File without changes
comfy/ldm/modules/diffusionmodules/model.py ADDED
@@ -0,0 +1,737 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ..attention import MemoryEfficientCrossAttention
10
+ from comfy import model_management
11
+ import comfy.ops
12
+
13
+ if model_management.xformers_enabled_vae():
14
+ import xformers
15
+ import xformers.ops
16
+
17
+ def get_timestep_embedding(timesteps, embedding_dim):
18
+ """
19
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
20
+ From Fairseq.
21
+ Build sinusoidal embeddings.
22
+ This matches the implementation in tensor2tensor, but differs slightly
23
+ from the description in Section 3.5 of "Attention Is All You Need".
24
+ """
25
+ assert len(timesteps.shape) == 1
26
+
27
+ half_dim = embedding_dim // 2
28
+ emb = math.log(10000) / (half_dim - 1)
29
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
30
+ emb = emb.to(device=timesteps.device)
31
+ emb = timesteps.float()[:, None] * emb[None, :]
32
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
33
+ if embedding_dim % 2 == 1: # zero pad
34
+ emb = torch.nn.functional.pad(emb, (0,1,0,0))
35
+ return emb
36
+
37
+
38
+ def nonlinearity(x):
39
+ # swish
40
+ return x*torch.sigmoid(x)
41
+
42
+
43
+ def Normalize(in_channels, num_groups=32):
44
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
45
+
46
+
47
+ class Upsample(nn.Module):
48
+ def __init__(self, in_channels, with_conv):
49
+ super().__init__()
50
+ self.with_conv = with_conv
51
+ if self.with_conv:
52
+ self.conv = comfy.ops.Conv2d(in_channels,
53
+ in_channels,
54
+ kernel_size=3,
55
+ stride=1,
56
+ padding=1)
57
+
58
+ def forward(self, x):
59
+ try:
60
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
61
+ except: #operation not implemented for bf16
62
+ b, c, h, w = x.shape
63
+ out = torch.empty((b, c, h*2, w*2), dtype=x.dtype, layout=x.layout, device=x.device)
64
+ split = 8
65
+ l = out.shape[1] // split
66
+ for i in range(0, out.shape[1], l):
67
+ out[:,i:i+l] = torch.nn.functional.interpolate(x[:,i:i+l].to(torch.float32), scale_factor=2.0, mode="nearest").to(x.dtype)
68
+ del x
69
+ x = out
70
+
71
+ if self.with_conv:
72
+ x = self.conv(x)
73
+ return x
74
+
75
+
76
+ class Downsample(nn.Module):
77
+ def __init__(self, in_channels, with_conv):
78
+ super().__init__()
79
+ self.with_conv = with_conv
80
+ if self.with_conv:
81
+ # no asymmetric padding in torch conv, must do it ourselves
82
+ self.conv = comfy.ops.Conv2d(in_channels,
83
+ in_channels,
84
+ kernel_size=3,
85
+ stride=2,
86
+ padding=0)
87
+
88
+ def forward(self, x):
89
+ if self.with_conv:
90
+ pad = (0,1,0,1)
91
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
92
+ x = self.conv(x)
93
+ else:
94
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
95
+ return x
96
+
97
+
98
+ class ResnetBlock(nn.Module):
99
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
100
+ dropout, temb_channels=512):
101
+ super().__init__()
102
+ self.in_channels = in_channels
103
+ out_channels = in_channels if out_channels is None else out_channels
104
+ self.out_channels = out_channels
105
+ self.use_conv_shortcut = conv_shortcut
106
+
107
+ self.swish = torch.nn.SiLU(inplace=True)
108
+ self.norm1 = Normalize(in_channels)
109
+ self.conv1 = comfy.ops.Conv2d(in_channels,
110
+ out_channels,
111
+ kernel_size=3,
112
+ stride=1,
113
+ padding=1)
114
+ if temb_channels > 0:
115
+ self.temb_proj = comfy.ops.Linear(temb_channels,
116
+ out_channels)
117
+ self.norm2 = Normalize(out_channels)
118
+ self.dropout = torch.nn.Dropout(dropout, inplace=True)
119
+ self.conv2 = comfy.ops.Conv2d(out_channels,
120
+ out_channels,
121
+ kernel_size=3,
122
+ stride=1,
123
+ padding=1)
124
+ if self.in_channels != self.out_channels:
125
+ if self.use_conv_shortcut:
126
+ self.conv_shortcut = comfy.ops.Conv2d(in_channels,
127
+ out_channels,
128
+ kernel_size=3,
129
+ stride=1,
130
+ padding=1)
131
+ else:
132
+ self.nin_shortcut = comfy.ops.Conv2d(in_channels,
133
+ out_channels,
134
+ kernel_size=1,
135
+ stride=1,
136
+ padding=0)
137
+
138
+ def forward(self, x, temb):
139
+ h = x
140
+ h = self.norm1(h)
141
+ h = self.swish(h)
142
+ h = self.conv1(h)
143
+
144
+ if temb is not None:
145
+ h = h + self.temb_proj(self.swish(temb))[:,:,None,None]
146
+
147
+ h = self.norm2(h)
148
+ h = self.swish(h)
149
+ h = self.dropout(h)
150
+ h = self.conv2(h)
151
+
152
+ if self.in_channels != self.out_channels:
153
+ if self.use_conv_shortcut:
154
+ x = self.conv_shortcut(x)
155
+ else:
156
+ x = self.nin_shortcut(x)
157
+
158
+ return x+h
159
+
160
+ def slice_attention(q, k, v):
161
+ r1 = torch.zeros_like(k, device=q.device)
162
+ scale = (int(q.shape[-1])**(-0.5))
163
+
164
+ mem_free_total = model_management.get_free_memory(q.device)
165
+
166
+ gb = 1024 ** 3
167
+ tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
168
+ modifier = 3 if q.element_size() == 2 else 2.5
169
+ mem_required = tensor_size * modifier
170
+ steps = 1
171
+
172
+ if mem_required > mem_free_total:
173
+ steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
174
+
175
+ while True:
176
+ try:
177
+ slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
178
+ for i in range(0, q.shape[1], slice_size):
179
+ end = i + slice_size
180
+ s1 = torch.bmm(q[:, i:end], k) * scale
181
+
182
+ s2 = torch.nn.functional.softmax(s1, dim=2).permute(0,2,1)
183
+ del s1
184
+
185
+ r1[:, :, i:end] = torch.bmm(v, s2)
186
+ del s2
187
+ break
188
+ except model_management.OOM_EXCEPTION as e:
189
+ model_management.soft_empty_cache(True)
190
+ steps *= 2
191
+ if steps > 128:
192
+ raise e
193
+ print("out of memory error, increasing steps and trying again", steps)
194
+
195
+ return r1
196
+
197
+ class AttnBlock(nn.Module):
198
+ def __init__(self, in_channels):
199
+ super().__init__()
200
+ self.in_channels = in_channels
201
+
202
+ self.norm = Normalize(in_channels)
203
+ self.q = comfy.ops.Conv2d(in_channels,
204
+ in_channels,
205
+ kernel_size=1,
206
+ stride=1,
207
+ padding=0)
208
+ self.k = comfy.ops.Conv2d(in_channels,
209
+ in_channels,
210
+ kernel_size=1,
211
+ stride=1,
212
+ padding=0)
213
+ self.v = comfy.ops.Conv2d(in_channels,
214
+ in_channels,
215
+ kernel_size=1,
216
+ stride=1,
217
+ padding=0)
218
+ self.proj_out = comfy.ops.Conv2d(in_channels,
219
+ in_channels,
220
+ kernel_size=1,
221
+ stride=1,
222
+ padding=0)
223
+
224
+ def forward(self, x):
225
+ h_ = x
226
+ h_ = self.norm(h_)
227
+ q = self.q(h_)
228
+ k = self.k(h_)
229
+ v = self.v(h_)
230
+
231
+ # compute attention
232
+ b,c,h,w = q.shape
233
+
234
+ q = q.reshape(b,c,h*w)
235
+ q = q.permute(0,2,1) # b,hw,c
236
+ k = k.reshape(b,c,h*w) # b,c,hw
237
+ v = v.reshape(b,c,h*w)
238
+
239
+ r1 = slice_attention(q, k, v)
240
+ h_ = r1.reshape(b,c,h,w)
241
+ del r1
242
+ h_ = self.proj_out(h_)
243
+
244
+ return x+h_
245
+
246
+ class MemoryEfficientAttnBlock(nn.Module):
247
+ """
248
+ Uses xformers efficient implementation,
249
+ see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
250
+ Note: this is a single-head self-attention operation
251
+ """
252
+ #
253
+ def __init__(self, in_channels):
254
+ super().__init__()
255
+ self.in_channels = in_channels
256
+
257
+ self.norm = Normalize(in_channels)
258
+ self.q = comfy.ops.Conv2d(in_channels,
259
+ in_channels,
260
+ kernel_size=1,
261
+ stride=1,
262
+ padding=0)
263
+ self.k = comfy.ops.Conv2d(in_channels,
264
+ in_channels,
265
+ kernel_size=1,
266
+ stride=1,
267
+ padding=0)
268
+ self.v = comfy.ops.Conv2d(in_channels,
269
+ in_channels,
270
+ kernel_size=1,
271
+ stride=1,
272
+ padding=0)
273
+ self.proj_out = comfy.ops.Conv2d(in_channels,
274
+ in_channels,
275
+ kernel_size=1,
276
+ stride=1,
277
+ padding=0)
278
+ self.attention_op: Optional[Any] = None
279
+
280
+ def forward(self, x):
281
+ h_ = x
282
+ h_ = self.norm(h_)
283
+ q = self.q(h_)
284
+ k = self.k(h_)
285
+ v = self.v(h_)
286
+
287
+ # compute attention
288
+ B, C, H, W = q.shape
289
+ q, k, v = map(
290
+ lambda t: t.view(B, C, -1).transpose(1, 2).contiguous(),
291
+ (q, k, v),
292
+ )
293
+
294
+ try:
295
+ out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
296
+ out = out.transpose(1, 2).reshape(B, C, H, W)
297
+ except NotImplementedError as e:
298
+ out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
299
+
300
+ out = self.proj_out(out)
301
+ return x+out
302
+
303
+ class MemoryEfficientAttnBlockPytorch(nn.Module):
304
+ def __init__(self, in_channels):
305
+ super().__init__()
306
+ self.in_channels = in_channels
307
+
308
+ self.norm = Normalize(in_channels)
309
+ self.q = comfy.ops.Conv2d(in_channels,
310
+ in_channels,
311
+ kernel_size=1,
312
+ stride=1,
313
+ padding=0)
314
+ self.k = comfy.ops.Conv2d(in_channels,
315
+ in_channels,
316
+ kernel_size=1,
317
+ stride=1,
318
+ padding=0)
319
+ self.v = comfy.ops.Conv2d(in_channels,
320
+ in_channels,
321
+ kernel_size=1,
322
+ stride=1,
323
+ padding=0)
324
+ self.proj_out = comfy.ops.Conv2d(in_channels,
325
+ in_channels,
326
+ kernel_size=1,
327
+ stride=1,
328
+ padding=0)
329
+ self.attention_op: Optional[Any] = None
330
+
331
+ def forward(self, x):
332
+ h_ = x
333
+ h_ = self.norm(h_)
334
+ q = self.q(h_)
335
+ k = self.k(h_)
336
+ v = self.v(h_)
337
+
338
+ # compute attention
339
+ B, C, H, W = q.shape
340
+ q, k, v = map(
341
+ lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(),
342
+ (q, k, v),
343
+ )
344
+
345
+ try:
346
+ out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
347
+ out = out.transpose(2, 3).reshape(B, C, H, W)
348
+ except model_management.OOM_EXCEPTION as e:
349
+ print("scaled_dot_product_attention OOMed: switched to slice attention")
350
+ out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
351
+
352
+ out = self.proj_out(out)
353
+ return x+out
354
+
355
+ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
356
+ def forward(self, x, context=None, mask=None):
357
+ b, c, h, w = x.shape
358
+ x = rearrange(x, 'b c h w -> b (h w) c')
359
+ out = super().forward(x, context=context, mask=mask)
360
+ out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w, c=c)
361
+ return x + out
362
+
363
+
364
+ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
365
+ assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown'
366
+ if model_management.xformers_enabled_vae() and attn_type == "vanilla":
367
+ attn_type = "vanilla-xformers"
368
+ if model_management.pytorch_attention_enabled() and attn_type == "vanilla":
369
+ attn_type = "vanilla-pytorch"
370
+ print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
371
+ if attn_type == "vanilla":
372
+ assert attn_kwargs is None
373
+ return AttnBlock(in_channels)
374
+ elif attn_type == "vanilla-xformers":
375
+ print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
376
+ return MemoryEfficientAttnBlock(in_channels)
377
+ elif attn_type == "vanilla-pytorch":
378
+ return MemoryEfficientAttnBlockPytorch(in_channels)
379
+ elif type == "memory-efficient-cross-attn":
380
+ attn_kwargs["query_dim"] = in_channels
381
+ return MemoryEfficientCrossAttentionWrapper(**attn_kwargs)
382
+ elif attn_type == "none":
383
+ return nn.Identity(in_channels)
384
+ else:
385
+ raise NotImplementedError()
386
+
387
+
388
+ class Model(nn.Module):
389
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
390
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
391
+ resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
392
+ super().__init__()
393
+ if use_linear_attn: attn_type = "linear"
394
+ self.ch = ch
395
+ self.temb_ch = self.ch*4
396
+ self.num_resolutions = len(ch_mult)
397
+ self.num_res_blocks = num_res_blocks
398
+ self.resolution = resolution
399
+ self.in_channels = in_channels
400
+
401
+ self.use_timestep = use_timestep
402
+ if self.use_timestep:
403
+ # timestep embedding
404
+ self.temb = nn.Module()
405
+ self.temb.dense = nn.ModuleList([
406
+ comfy.ops.Linear(self.ch,
407
+ self.temb_ch),
408
+ comfy.ops.Linear(self.temb_ch,
409
+ self.temb_ch),
410
+ ])
411
+
412
+ # downsampling
413
+ self.conv_in = comfy.ops.Conv2d(in_channels,
414
+ self.ch,
415
+ kernel_size=3,
416
+ stride=1,
417
+ padding=1)
418
+
419
+ curr_res = resolution
420
+ in_ch_mult = (1,)+tuple(ch_mult)
421
+ self.down = nn.ModuleList()
422
+ for i_level in range(self.num_resolutions):
423
+ block = nn.ModuleList()
424
+ attn = nn.ModuleList()
425
+ block_in = ch*in_ch_mult[i_level]
426
+ block_out = ch*ch_mult[i_level]
427
+ for i_block in range(self.num_res_blocks):
428
+ block.append(ResnetBlock(in_channels=block_in,
429
+ out_channels=block_out,
430
+ temb_channels=self.temb_ch,
431
+ dropout=dropout))
432
+ block_in = block_out
433
+ if curr_res in attn_resolutions:
434
+ attn.append(make_attn(block_in, attn_type=attn_type))
435
+ down = nn.Module()
436
+ down.block = block
437
+ down.attn = attn
438
+ if i_level != self.num_resolutions-1:
439
+ down.downsample = Downsample(block_in, resamp_with_conv)
440
+ curr_res = curr_res // 2
441
+ self.down.append(down)
442
+
443
+ # middle
444
+ self.mid = nn.Module()
445
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
446
+ out_channels=block_in,
447
+ temb_channels=self.temb_ch,
448
+ dropout=dropout)
449
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
450
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
451
+ out_channels=block_in,
452
+ temb_channels=self.temb_ch,
453
+ dropout=dropout)
454
+
455
+ # upsampling
456
+ self.up = nn.ModuleList()
457
+ for i_level in reversed(range(self.num_resolutions)):
458
+ block = nn.ModuleList()
459
+ attn = nn.ModuleList()
460
+ block_out = ch*ch_mult[i_level]
461
+ skip_in = ch*ch_mult[i_level]
462
+ for i_block in range(self.num_res_blocks+1):
463
+ if i_block == self.num_res_blocks:
464
+ skip_in = ch*in_ch_mult[i_level]
465
+ block.append(ResnetBlock(in_channels=block_in+skip_in,
466
+ out_channels=block_out,
467
+ temb_channels=self.temb_ch,
468
+ dropout=dropout))
469
+ block_in = block_out
470
+ if curr_res in attn_resolutions:
471
+ attn.append(make_attn(block_in, attn_type=attn_type))
472
+ up = nn.Module()
473
+ up.block = block
474
+ up.attn = attn
475
+ if i_level != 0:
476
+ up.upsample = Upsample(block_in, resamp_with_conv)
477
+ curr_res = curr_res * 2
478
+ self.up.insert(0, up) # prepend to get consistent order
479
+
480
+ # end
481
+ self.norm_out = Normalize(block_in)
482
+ self.conv_out = comfy.ops.Conv2d(block_in,
483
+ out_ch,
484
+ kernel_size=3,
485
+ stride=1,
486
+ padding=1)
487
+
488
+ def forward(self, x, t=None, context=None):
489
+ #assert x.shape[2] == x.shape[3] == self.resolution
490
+ if context is not None:
491
+ # assume aligned context, cat along channel axis
492
+ x = torch.cat((x, context), dim=1)
493
+ if self.use_timestep:
494
+ # timestep embedding
495
+ assert t is not None
496
+ temb = get_timestep_embedding(t, self.ch)
497
+ temb = self.temb.dense[0](temb)
498
+ temb = nonlinearity(temb)
499
+ temb = self.temb.dense[1](temb)
500
+ else:
501
+ temb = None
502
+
503
+ # downsampling
504
+ hs = [self.conv_in(x)]
505
+ for i_level in range(self.num_resolutions):
506
+ for i_block in range(self.num_res_blocks):
507
+ h = self.down[i_level].block[i_block](hs[-1], temb)
508
+ if len(self.down[i_level].attn) > 0:
509
+ h = self.down[i_level].attn[i_block](h)
510
+ hs.append(h)
511
+ if i_level != self.num_resolutions-1:
512
+ hs.append(self.down[i_level].downsample(hs[-1]))
513
+
514
+ # middle
515
+ h = hs[-1]
516
+ h = self.mid.block_1(h, temb)
517
+ h = self.mid.attn_1(h)
518
+ h = self.mid.block_2(h, temb)
519
+
520
+ # upsampling
521
+ for i_level in reversed(range(self.num_resolutions)):
522
+ for i_block in range(self.num_res_blocks+1):
523
+ h = self.up[i_level].block[i_block](
524
+ torch.cat([h, hs.pop()], dim=1), temb)
525
+ if len(self.up[i_level].attn) > 0:
526
+ h = self.up[i_level].attn[i_block](h)
527
+ if i_level != 0:
528
+ h = self.up[i_level].upsample(h)
529
+
530
+ # end
531
+ h = self.norm_out(h)
532
+ h = nonlinearity(h)
533
+ h = self.conv_out(h)
534
+ return h
535
+
536
+ def get_last_layer(self):
537
+ return self.conv_out.weight
538
+
539
+
540
+ class Encoder(nn.Module):
541
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
542
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
543
+ resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
544
+ **ignore_kwargs):
545
+ super().__init__()
546
+ if use_linear_attn: attn_type = "linear"
547
+ self.ch = ch
548
+ self.temb_ch = 0
549
+ self.num_resolutions = len(ch_mult)
550
+ self.num_res_blocks = num_res_blocks
551
+ self.resolution = resolution
552
+ self.in_channels = in_channels
553
+
554
+ # downsampling
555
+ self.conv_in = comfy.ops.Conv2d(in_channels,
556
+ self.ch,
557
+ kernel_size=3,
558
+ stride=1,
559
+ padding=1)
560
+
561
+ curr_res = resolution
562
+ in_ch_mult = (1,)+tuple(ch_mult)
563
+ self.in_ch_mult = in_ch_mult
564
+ self.down = nn.ModuleList()
565
+ for i_level in range(self.num_resolutions):
566
+ block = nn.ModuleList()
567
+ attn = nn.ModuleList()
568
+ block_in = ch*in_ch_mult[i_level]
569
+ block_out = ch*ch_mult[i_level]
570
+ for i_block in range(self.num_res_blocks):
571
+ block.append(ResnetBlock(in_channels=block_in,
572
+ out_channels=block_out,
573
+ temb_channels=self.temb_ch,
574
+ dropout=dropout))
575
+ block_in = block_out
576
+ if curr_res in attn_resolutions:
577
+ attn.append(make_attn(block_in, attn_type=attn_type))
578
+ down = nn.Module()
579
+ down.block = block
580
+ down.attn = attn
581
+ if i_level != self.num_resolutions-1:
582
+ down.downsample = Downsample(block_in, resamp_with_conv)
583
+ curr_res = curr_res // 2
584
+ self.down.append(down)
585
+
586
+ # middle
587
+ self.mid = nn.Module()
588
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
589
+ out_channels=block_in,
590
+ temb_channels=self.temb_ch,
591
+ dropout=dropout)
592
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
593
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
594
+ out_channels=block_in,
595
+ temb_channels=self.temb_ch,
596
+ dropout=dropout)
597
+
598
+ # end
599
+ self.norm_out = Normalize(block_in)
600
+ self.conv_out = comfy.ops.Conv2d(block_in,
601
+ 2*z_channels if double_z else z_channels,
602
+ kernel_size=3,
603
+ stride=1,
604
+ padding=1)
605
+
606
+ def forward(self, x):
607
+ # timestep embedding
608
+ temb = None
609
+ # downsampling
610
+ h = self.conv_in(x)
611
+ for i_level in range(self.num_resolutions):
612
+ for i_block in range(self.num_res_blocks):
613
+ h = self.down[i_level].block[i_block](h, temb)
614
+ if len(self.down[i_level].attn) > 0:
615
+ h = self.down[i_level].attn[i_block](h)
616
+ if i_level != self.num_resolutions-1:
617
+ h = self.down[i_level].downsample(h)
618
+
619
+ # middle
620
+ h = self.mid.block_1(h, temb)
621
+ h = self.mid.attn_1(h)
622
+ h = self.mid.block_2(h, temb)
623
+
624
+ # end
625
+ h = self.norm_out(h)
626
+ h = nonlinearity(h)
627
+ h = self.conv_out(h)
628
+ return h
629
+
630
+
631
+ class Decoder(nn.Module):
632
+ def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
633
+ attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
634
+ resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
635
+ attn_type="vanilla", **ignorekwargs):
636
+ super().__init__()
637
+ if use_linear_attn: attn_type = "linear"
638
+ self.ch = ch
639
+ self.temb_ch = 0
640
+ self.num_resolutions = len(ch_mult)
641
+ self.num_res_blocks = num_res_blocks
642
+ self.resolution = resolution
643
+ self.in_channels = in_channels
644
+ self.give_pre_end = give_pre_end
645
+ self.tanh_out = tanh_out
646
+
647
+ # compute in_ch_mult, block_in and curr_res at lowest res
648
+ in_ch_mult = (1,)+tuple(ch_mult)
649
+ block_in = ch*ch_mult[self.num_resolutions-1]
650
+ curr_res = resolution // 2**(self.num_resolutions-1)
651
+ self.z_shape = (1,z_channels,curr_res,curr_res)
652
+ print("Working with z of shape {} = {} dimensions.".format(
653
+ self.z_shape, np.prod(self.z_shape)))
654
+
655
+ # z to block_in
656
+ self.conv_in = comfy.ops.Conv2d(z_channels,
657
+ block_in,
658
+ kernel_size=3,
659
+ stride=1,
660
+ padding=1)
661
+
662
+ # middle
663
+ self.mid = nn.Module()
664
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
665
+ out_channels=block_in,
666
+ temb_channels=self.temb_ch,
667
+ dropout=dropout)
668
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
669
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
670
+ out_channels=block_in,
671
+ temb_channels=self.temb_ch,
672
+ dropout=dropout)
673
+
674
+ # upsampling
675
+ self.up = nn.ModuleList()
676
+ for i_level in reversed(range(self.num_resolutions)):
677
+ block = nn.ModuleList()
678
+ attn = nn.ModuleList()
679
+ block_out = ch*ch_mult[i_level]
680
+ for i_block in range(self.num_res_blocks+1):
681
+ block.append(ResnetBlock(in_channels=block_in,
682
+ out_channels=block_out,
683
+ temb_channels=self.temb_ch,
684
+ dropout=dropout))
685
+ block_in = block_out
686
+ if curr_res in attn_resolutions:
687
+ attn.append(make_attn(block_in, attn_type=attn_type))
688
+ up = nn.Module()
689
+ up.block = block
690
+ up.attn = attn
691
+ if i_level != 0:
692
+ up.upsample = Upsample(block_in, resamp_with_conv)
693
+ curr_res = curr_res * 2
694
+ self.up.insert(0, up) # prepend to get consistent order
695
+
696
+ # end
697
+ self.norm_out = Normalize(block_in)
698
+ self.conv_out = comfy.ops.Conv2d(block_in,
699
+ out_ch,
700
+ kernel_size=3,
701
+ stride=1,
702
+ padding=1)
703
+
704
+ def forward(self, z):
705
+ #assert z.shape[1:] == self.z_shape[1:]
706
+ self.last_z_shape = z.shape
707
+
708
+ # timestep embedding
709
+ temb = None
710
+
711
+ # z to block_in
712
+ h = self.conv_in(z)
713
+
714
+ # middle
715
+ h = self.mid.block_1(h, temb)
716
+ h = self.mid.attn_1(h)
717
+ h = self.mid.block_2(h, temb)
718
+
719
+ # upsampling
720
+ for i_level in reversed(range(self.num_resolutions)):
721
+ for i_block in range(self.num_res_blocks+1):
722
+ h = self.up[i_level].block[i_block](h, temb)
723
+ if len(self.up[i_level].attn) > 0:
724
+ h = self.up[i_level].attn[i_block](h)
725
+ if i_level != 0:
726
+ h = self.up[i_level].upsample(h)
727
+
728
+ # end
729
+ if self.give_pre_end:
730
+ return h
731
+
732
+ h = self.norm_out(h)
733
+ h = nonlinearity(h)
734
+ h = self.conv_out(h)
735
+ if self.tanh_out:
736
+ h = torch.tanh(h)
737
+ return h
comfy/ldm/modules/diffusionmodules/openaimodel.py ADDED
@@ -0,0 +1,664 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 .util import (
10
+ checkpoint,
11
+ avg_pool_nd,
12
+ zero_module,
13
+ normalization,
14
+ timestep_embedding,
15
+ )
16
+ from ..attention import SpatialTransformer
17
+ from comfy.ldm.util import exists
18
+ import comfy.ops
19
+
20
+ class TimestepBlock(nn.Module):
21
+ """
22
+ Any module where forward() takes timestep embeddings as a second argument.
23
+ """
24
+
25
+ @abstractmethod
26
+ def forward(self, x, emb):
27
+ """
28
+ Apply the module to `x` given `emb` timestep embeddings.
29
+ """
30
+
31
+
32
+ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
33
+ """
34
+ A sequential module that passes timestep embeddings to the children that
35
+ support it as an extra input.
36
+ """
37
+
38
+ def forward(self, x, emb, context=None, transformer_options={}, output_shape=None):
39
+ for layer in self:
40
+ if isinstance(layer, TimestepBlock):
41
+ x = layer(x, emb)
42
+ elif isinstance(layer, SpatialTransformer):
43
+ x = layer(x, context, transformer_options)
44
+ elif isinstance(layer, Upsample):
45
+ x = layer(x, output_shape=output_shape)
46
+ else:
47
+ x = layer(x)
48
+ return x
49
+
50
+ #This is needed because accelerate makes a copy of transformer_options which breaks "current_index"
51
+ def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None):
52
+ for layer in ts:
53
+ if isinstance(layer, TimestepBlock):
54
+ x = layer(x, emb)
55
+ elif isinstance(layer, SpatialTransformer):
56
+ x = layer(x, context, transformer_options)
57
+ transformer_options["current_index"] += 1
58
+ elif isinstance(layer, Upsample):
59
+ x = layer(x, output_shape=output_shape)
60
+ else:
61
+ x = layer(x)
62
+ return x
63
+
64
+ class Upsample(nn.Module):
65
+ """
66
+ An upsampling layer with an optional convolution.
67
+ :param channels: channels in the inputs and outputs.
68
+ :param use_conv: a bool determining if a convolution is applied.
69
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
70
+ upsampling occurs in the inner-two dimensions.
71
+ """
72
+
73
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=comfy.ops):
74
+ super().__init__()
75
+ self.channels = channels
76
+ self.out_channels = out_channels or channels
77
+ self.use_conv = use_conv
78
+ self.dims = dims
79
+ if use_conv:
80
+ self.conv = operations.conv_nd(dims, self.channels, self.out_channels, 3, padding=padding, dtype=dtype, device=device)
81
+
82
+ def forward(self, x, output_shape=None):
83
+ assert x.shape[1] == self.channels
84
+ if self.dims == 3:
85
+ shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
86
+ if output_shape is not None:
87
+ shape[1] = output_shape[3]
88
+ shape[2] = output_shape[4]
89
+ else:
90
+ shape = [x.shape[2] * 2, x.shape[3] * 2]
91
+ if output_shape is not None:
92
+ shape[0] = output_shape[2]
93
+ shape[1] = output_shape[3]
94
+
95
+ x = F.interpolate(x, size=shape, mode="nearest")
96
+ if self.use_conv:
97
+ x = self.conv(x)
98
+ return x
99
+
100
+ class Downsample(nn.Module):
101
+ """
102
+ A downsampling layer with an optional convolution.
103
+ :param channels: channels in the inputs and outputs.
104
+ :param use_conv: a bool determining if a convolution is applied.
105
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
106
+ downsampling occurs in the inner-two dimensions.
107
+ """
108
+
109
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=comfy.ops):
110
+ super().__init__()
111
+ self.channels = channels
112
+ self.out_channels = out_channels or channels
113
+ self.use_conv = use_conv
114
+ self.dims = dims
115
+ stride = 2 if dims != 3 else (1, 2, 2)
116
+ if use_conv:
117
+ self.op = operations.conv_nd(
118
+ dims, self.channels, self.out_channels, 3, stride=stride, padding=padding, dtype=dtype, device=device
119
+ )
120
+ else:
121
+ assert self.channels == self.out_channels
122
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
123
+
124
+ def forward(self, x):
125
+ assert x.shape[1] == self.channels
126
+ return self.op(x)
127
+
128
+
129
+ class ResBlock(TimestepBlock):
130
+ """
131
+ A residual block that can optionally change the number of channels.
132
+ :param channels: the number of input channels.
133
+ :param emb_channels: the number of timestep embedding channels.
134
+ :param dropout: the rate of dropout.
135
+ :param out_channels: if specified, the number of out channels.
136
+ :param use_conv: if True and out_channels is specified, use a spatial
137
+ convolution instead of a smaller 1x1 convolution to change the
138
+ channels in the skip connection.
139
+ :param dims: determines if the signal is 1D, 2D, or 3D.
140
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
141
+ :param up: if True, use this block for upsampling.
142
+ :param down: if True, use this block for downsampling.
143
+ """
144
+
145
+ def __init__(
146
+ self,
147
+ channels,
148
+ emb_channels,
149
+ dropout,
150
+ out_channels=None,
151
+ use_conv=False,
152
+ use_scale_shift_norm=False,
153
+ dims=2,
154
+ use_checkpoint=False,
155
+ up=False,
156
+ down=False,
157
+ dtype=None,
158
+ device=None,
159
+ operations=comfy.ops
160
+ ):
161
+ super().__init__()
162
+ self.channels = channels
163
+ self.emb_channels = emb_channels
164
+ self.dropout = dropout
165
+ self.out_channels = out_channels or channels
166
+ self.use_conv = use_conv
167
+ self.use_checkpoint = use_checkpoint
168
+ self.use_scale_shift_norm = use_scale_shift_norm
169
+
170
+ self.in_layers = nn.Sequential(
171
+ nn.GroupNorm(32, channels, dtype=dtype, device=device),
172
+ nn.SiLU(),
173
+ operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device),
174
+ )
175
+
176
+ self.updown = up or down
177
+
178
+ if up:
179
+ self.h_upd = Upsample(channels, False, dims, dtype=dtype, device=device)
180
+ self.x_upd = Upsample(channels, False, dims, dtype=dtype, device=device)
181
+ elif down:
182
+ self.h_upd = Downsample(channels, False, dims, dtype=dtype, device=device)
183
+ self.x_upd = Downsample(channels, False, dims, dtype=dtype, device=device)
184
+ else:
185
+ self.h_upd = self.x_upd = nn.Identity()
186
+
187
+ self.emb_layers = nn.Sequential(
188
+ nn.SiLU(),
189
+ operations.Linear(
190
+ emb_channels,
191
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels, dtype=dtype, device=device
192
+ ),
193
+ )
194
+ self.out_layers = nn.Sequential(
195
+ nn.GroupNorm(32, self.out_channels, dtype=dtype, device=device),
196
+ nn.SiLU(),
197
+ nn.Dropout(p=dropout),
198
+ zero_module(
199
+ operations.conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, dtype=dtype, device=device)
200
+ ),
201
+ )
202
+
203
+ if self.out_channels == channels:
204
+ self.skip_connection = nn.Identity()
205
+ elif use_conv:
206
+ self.skip_connection = operations.conv_nd(
207
+ dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device
208
+ )
209
+ else:
210
+ self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 1, dtype=dtype, device=device)
211
+
212
+ def forward(self, x, emb):
213
+ """
214
+ Apply the block to a Tensor, conditioned on a timestep embedding.
215
+ :param x: an [N x C x ...] Tensor of features.
216
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
217
+ :return: an [N x C x ...] Tensor of outputs.
218
+ """
219
+ return checkpoint(
220
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
221
+ )
222
+
223
+
224
+ def _forward(self, x, emb):
225
+ if self.updown:
226
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
227
+ h = in_rest(x)
228
+ h = self.h_upd(h)
229
+ x = self.x_upd(x)
230
+ h = in_conv(h)
231
+ else:
232
+ h = self.in_layers(x)
233
+ emb_out = self.emb_layers(emb).type(h.dtype)
234
+ while len(emb_out.shape) < len(h.shape):
235
+ emb_out = emb_out[..., None]
236
+ if self.use_scale_shift_norm:
237
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
238
+ scale, shift = th.chunk(emb_out, 2, dim=1)
239
+ h = out_norm(h) * (1 + scale) + shift
240
+ h = out_rest(h)
241
+ else:
242
+ h = h + emb_out
243
+ h = self.out_layers(h)
244
+ return self.skip_connection(x) + h
245
+
246
+ class Timestep(nn.Module):
247
+ def __init__(self, dim):
248
+ super().__init__()
249
+ self.dim = dim
250
+
251
+ def forward(self, t):
252
+ return timestep_embedding(t, self.dim)
253
+
254
+
255
+ class UNetModel(nn.Module):
256
+ """
257
+ The full UNet model with attention and timestep embedding.
258
+ :param in_channels: channels in the input Tensor.
259
+ :param model_channels: base channel count for the model.
260
+ :param out_channels: channels in the output Tensor.
261
+ :param num_res_blocks: number of residual blocks per downsample.
262
+ :param attention_resolutions: a collection of downsample rates at which
263
+ attention will take place. May be a set, list, or tuple.
264
+ For example, if this contains 4, then at 4x downsampling, attention
265
+ will be used.
266
+ :param dropout: the dropout probability.
267
+ :param channel_mult: channel multiplier for each level of the UNet.
268
+ :param conv_resample: if True, use learned convolutions for upsampling and
269
+ downsampling.
270
+ :param dims: determines if the signal is 1D, 2D, or 3D.
271
+ :param num_classes: if specified (as an int), then this model will be
272
+ class-conditional with `num_classes` classes.
273
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
274
+ :param num_heads: the number of attention heads in each attention layer.
275
+ :param num_heads_channels: if specified, ignore num_heads and instead use
276
+ a fixed channel width per attention head.
277
+ :param num_heads_upsample: works with num_heads to set a different number
278
+ of heads for upsampling. Deprecated.
279
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
280
+ :param resblock_updown: use residual blocks for up/downsampling.
281
+ :param use_new_attention_order: use a different attention pattern for potentially
282
+ increased efficiency.
283
+ """
284
+
285
+ def __init__(
286
+ self,
287
+ image_size,
288
+ in_channels,
289
+ model_channels,
290
+ out_channels,
291
+ num_res_blocks,
292
+ attention_resolutions,
293
+ dropout=0,
294
+ channel_mult=(1, 2, 4, 8),
295
+ conv_resample=True,
296
+ dims=2,
297
+ num_classes=None,
298
+ use_checkpoint=False,
299
+ use_fp16=False,
300
+ use_bf16=False,
301
+ num_heads=-1,
302
+ num_head_channels=-1,
303
+ num_heads_upsample=-1,
304
+ use_scale_shift_norm=False,
305
+ resblock_updown=False,
306
+ use_new_attention_order=False,
307
+ use_spatial_transformer=False, # custom transformer support
308
+ transformer_depth=1, # custom transformer support
309
+ context_dim=None, # custom transformer support
310
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
311
+ legacy=True,
312
+ disable_self_attentions=None,
313
+ num_attention_blocks=None,
314
+ disable_middle_self_attn=False,
315
+ use_linear_in_transformer=False,
316
+ adm_in_channels=None,
317
+ transformer_depth_middle=None,
318
+ device=None,
319
+ operations=comfy.ops,
320
+ ):
321
+ super().__init__()
322
+ assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
323
+ if use_spatial_transformer:
324
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
325
+
326
+ if context_dim is not None:
327
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
328
+ # from omegaconf.listconfig import ListConfig
329
+ # if type(context_dim) == ListConfig:
330
+ # context_dim = list(context_dim)
331
+
332
+ if num_heads_upsample == -1:
333
+ num_heads_upsample = num_heads
334
+
335
+ if num_heads == -1:
336
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
337
+
338
+ if num_head_channels == -1:
339
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
340
+
341
+ self.image_size = image_size
342
+ self.in_channels = in_channels
343
+ self.model_channels = model_channels
344
+ self.out_channels = out_channels
345
+ if isinstance(transformer_depth, int):
346
+ transformer_depth = len(channel_mult) * [transformer_depth]
347
+ if transformer_depth_middle is None:
348
+ transformer_depth_middle = transformer_depth[-1]
349
+ if isinstance(num_res_blocks, int):
350
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
351
+ else:
352
+ if len(num_res_blocks) != len(channel_mult):
353
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
354
+ "as a list/tuple (per-level) with the same length as channel_mult")
355
+ self.num_res_blocks = num_res_blocks
356
+ if disable_self_attentions is not None:
357
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
358
+ assert len(disable_self_attentions) == len(channel_mult)
359
+ if num_attention_blocks is not None:
360
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
361
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
362
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
363
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
364
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
365
+ f"attention will still not be set.")
366
+
367
+ self.attention_resolutions = attention_resolutions
368
+ self.dropout = dropout
369
+ self.channel_mult = channel_mult
370
+ self.conv_resample = conv_resample
371
+ self.num_classes = num_classes
372
+ self.use_checkpoint = use_checkpoint
373
+ self.dtype = th.float16 if use_fp16 else th.float32
374
+ self.dtype = th.bfloat16 if use_bf16 else self.dtype
375
+ self.num_heads = num_heads
376
+ self.num_head_channels = num_head_channels
377
+ self.num_heads_upsample = num_heads_upsample
378
+ self.predict_codebook_ids = n_embed is not None
379
+
380
+ time_embed_dim = model_channels * 4
381
+ self.time_embed = nn.Sequential(
382
+ operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
383
+ nn.SiLU(),
384
+ operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
385
+ )
386
+
387
+ if self.num_classes is not None:
388
+ if isinstance(self.num_classes, int):
389
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
390
+ elif self.num_classes == "continuous":
391
+ print("setting up linear c_adm embedding layer")
392
+ self.label_emb = nn.Linear(1, time_embed_dim)
393
+ elif self.num_classes == "sequential":
394
+ assert adm_in_channels is not None
395
+ self.label_emb = nn.Sequential(
396
+ nn.Sequential(
397
+ operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
398
+ nn.SiLU(),
399
+ operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
400
+ )
401
+ )
402
+ else:
403
+ raise ValueError()
404
+
405
+ self.input_blocks = nn.ModuleList(
406
+ [
407
+ TimestepEmbedSequential(
408
+ operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
409
+ )
410
+ ]
411
+ )
412
+ self._feature_size = model_channels
413
+ input_block_chans = [model_channels]
414
+ ch = model_channels
415
+ ds = 1
416
+ for level, mult in enumerate(channel_mult):
417
+ for nr in range(self.num_res_blocks[level]):
418
+ layers = [
419
+ ResBlock(
420
+ ch,
421
+ time_embed_dim,
422
+ dropout,
423
+ out_channels=mult * model_channels,
424
+ dims=dims,
425
+ use_checkpoint=use_checkpoint,
426
+ use_scale_shift_norm=use_scale_shift_norm,
427
+ dtype=self.dtype,
428
+ device=device,
429
+ operations=operations,
430
+ )
431
+ ]
432
+ ch = mult * model_channels
433
+ if ds in attention_resolutions:
434
+ if num_head_channels == -1:
435
+ dim_head = ch // num_heads
436
+ else:
437
+ num_heads = ch // num_head_channels
438
+ dim_head = num_head_channels
439
+ if legacy:
440
+ #num_heads = 1
441
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
442
+ if exists(disable_self_attentions):
443
+ disabled_sa = disable_self_attentions[level]
444
+ else:
445
+ disabled_sa = False
446
+
447
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
448
+ layers.append(SpatialTransformer(
449
+ ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim,
450
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
451
+ use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
452
+ )
453
+ )
454
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
455
+ self._feature_size += ch
456
+ input_block_chans.append(ch)
457
+ if level != len(channel_mult) - 1:
458
+ out_ch = ch
459
+ self.input_blocks.append(
460
+ TimestepEmbedSequential(
461
+ ResBlock(
462
+ ch,
463
+ time_embed_dim,
464
+ dropout,
465
+ out_channels=out_ch,
466
+ dims=dims,
467
+ use_checkpoint=use_checkpoint,
468
+ use_scale_shift_norm=use_scale_shift_norm,
469
+ down=True,
470
+ dtype=self.dtype,
471
+ device=device,
472
+ operations=operations
473
+ )
474
+ if resblock_updown
475
+ else Downsample(
476
+ ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
477
+ )
478
+ )
479
+ )
480
+ ch = out_ch
481
+ input_block_chans.append(ch)
482
+ ds *= 2
483
+ self._feature_size += ch
484
+
485
+ if num_head_channels == -1:
486
+ dim_head = ch // num_heads
487
+ else:
488
+ num_heads = ch // num_head_channels
489
+ dim_head = num_head_channels
490
+ if legacy:
491
+ #num_heads = 1
492
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
493
+ self.middle_block = TimestepEmbedSequential(
494
+ ResBlock(
495
+ ch,
496
+ time_embed_dim,
497
+ dropout,
498
+ dims=dims,
499
+ use_checkpoint=use_checkpoint,
500
+ use_scale_shift_norm=use_scale_shift_norm,
501
+ dtype=self.dtype,
502
+ device=device,
503
+ operations=operations
504
+ ),
505
+ SpatialTransformer( # always uses a self-attn
506
+ ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
507
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
508
+ use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
509
+ ),
510
+ ResBlock(
511
+ ch,
512
+ time_embed_dim,
513
+ dropout,
514
+ dims=dims,
515
+ use_checkpoint=use_checkpoint,
516
+ use_scale_shift_norm=use_scale_shift_norm,
517
+ dtype=self.dtype,
518
+ device=device,
519
+ operations=operations
520
+ ),
521
+ )
522
+ self._feature_size += ch
523
+
524
+ self.output_blocks = nn.ModuleList([])
525
+ for level, mult in list(enumerate(channel_mult))[::-1]:
526
+ for i in range(self.num_res_blocks[level] + 1):
527
+ ich = input_block_chans.pop()
528
+ layers = [
529
+ ResBlock(
530
+ ch + ich,
531
+ time_embed_dim,
532
+ dropout,
533
+ out_channels=model_channels * mult,
534
+ dims=dims,
535
+ use_checkpoint=use_checkpoint,
536
+ use_scale_shift_norm=use_scale_shift_norm,
537
+ dtype=self.dtype,
538
+ device=device,
539
+ operations=operations
540
+ )
541
+ ]
542
+ ch = model_channels * mult
543
+ if ds in attention_resolutions:
544
+ if num_head_channels == -1:
545
+ dim_head = ch // num_heads
546
+ else:
547
+ num_heads = ch // num_head_channels
548
+ dim_head = num_head_channels
549
+ if legacy:
550
+ #num_heads = 1
551
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
552
+ if exists(disable_self_attentions):
553
+ disabled_sa = disable_self_attentions[level]
554
+ else:
555
+ disabled_sa = False
556
+
557
+ if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
558
+ layers.append(
559
+ SpatialTransformer(
560
+ ch, num_heads, dim_head, depth=transformer_depth[level], context_dim=context_dim,
561
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
562
+ use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
563
+ )
564
+ )
565
+ if level and i == self.num_res_blocks[level]:
566
+ out_ch = ch
567
+ layers.append(
568
+ ResBlock(
569
+ ch,
570
+ time_embed_dim,
571
+ dropout,
572
+ out_channels=out_ch,
573
+ dims=dims,
574
+ use_checkpoint=use_checkpoint,
575
+ use_scale_shift_norm=use_scale_shift_norm,
576
+ up=True,
577
+ dtype=self.dtype,
578
+ device=device,
579
+ operations=operations
580
+ )
581
+ if resblock_updown
582
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations)
583
+ )
584
+ ds //= 2
585
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
586
+ self._feature_size += ch
587
+
588
+ self.out = nn.Sequential(
589
+ nn.GroupNorm(32, ch, dtype=self.dtype, device=device),
590
+ nn.SiLU(),
591
+ zero_module(operations.conv_nd(dims, model_channels, out_channels, 3, padding=1, dtype=self.dtype, device=device)),
592
+ )
593
+ if self.predict_codebook_ids:
594
+ self.id_predictor = nn.Sequential(
595
+ nn.GroupNorm(32, ch, dtype=self.dtype, device=device),
596
+ operations.conv_nd(dims, model_channels, n_embed, 1, dtype=self.dtype, device=device),
597
+ #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
598
+ )
599
+
600
+ def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
601
+ """
602
+ Apply the model to an input batch.
603
+ :param x: an [N x C x ...] Tensor of inputs.
604
+ :param timesteps: a 1-D batch of timesteps.
605
+ :param context: conditioning plugged in via crossattn
606
+ :param y: an [N] Tensor of labels, if class-conditional.
607
+ :return: an [N x C x ...] Tensor of outputs.
608
+ """
609
+ transformer_options["original_shape"] = list(x.shape)
610
+ transformer_options["current_index"] = 0
611
+ transformer_patches = transformer_options.get("patches", {})
612
+
613
+ assert (y is not None) == (
614
+ self.num_classes is not None
615
+ ), "must specify y if and only if the model is class-conditional"
616
+ hs = []
617
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(self.dtype)
618
+ emb = self.time_embed(t_emb)
619
+
620
+ if self.num_classes is not None:
621
+ assert y.shape[0] == x.shape[0]
622
+ emb = emb + self.label_emb(y)
623
+
624
+ h = x.type(self.dtype)
625
+ for id, module in enumerate(self.input_blocks):
626
+ transformer_options["block"] = ("input", id)
627
+ h = forward_timestep_embed(module, h, emb, context, transformer_options)
628
+ if control is not None and 'input' in control and len(control['input']) > 0:
629
+ ctrl = control['input'].pop()
630
+ if ctrl is not None:
631
+ h += ctrl
632
+ hs.append(h)
633
+ transformer_options["block"] = ("middle", 0)
634
+ h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options)
635
+ if control is not None and 'middle' in control and len(control['middle']) > 0:
636
+ ctrl = control['middle'].pop()
637
+ if ctrl is not None:
638
+ h += ctrl
639
+
640
+ for id, module in enumerate(self.output_blocks):
641
+ transformer_options["block"] = ("output", id)
642
+ hsp = hs.pop()
643
+ if control is not None and 'output' in control and len(control['output']) > 0:
644
+ ctrl = control['output'].pop()
645
+ if ctrl is not None:
646
+ hsp += ctrl
647
+
648
+ if "output_block_patch" in transformer_patches:
649
+ patch = transformer_patches["output_block_patch"]
650
+ for p in patch:
651
+ h, hsp = p(h, hsp, transformer_options)
652
+
653
+ h = th.cat([h, hsp], dim=1)
654
+ del hsp
655
+ if len(hs) > 0:
656
+ output_shape = hs[-1].shape
657
+ else:
658
+ output_shape = None
659
+ h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape)
660
+ h = h.type(x.dtype)
661
+ if self.predict_codebook_ids:
662
+ return self.id_predictor(h)
663
+ else:
664
+ 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 .util import extract_into_tensor, make_beta_schedule
7
+ from comfy.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 comfy.ldm.util import instantiate_from_config
19
+ import comfy.ops
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, dtype=None):
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, dtype=dtype)
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 comfy.ops.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 comfy.ops.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/noise_aug_modules.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ..diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
2
+ from ..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/sub_quadratic_attention.py ADDED
@@ -0,0 +1,250 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # original source:
2
+ # https://github.com/AminRezaei0x443/memory-efficient-attention/blob/1bc0d9e6ac5f82ea43a375135c4e1d3896ee1694/memory_efficient_attention/attention_torch.py
3
+ # license:
4
+ # MIT
5
+ # credit:
6
+ # Amin Rezaei (original author)
7
+ # Alex Birch (optimized algorithm for 3D tensors, at the expense of removing bias, masking and callbacks)
8
+ # implementation of:
9
+ # Self-attention Does Not Need O(n2) Memory":
10
+ # https://arxiv.org/abs/2112.05682v2
11
+
12
+ from functools import partial
13
+ import torch
14
+ from torch import Tensor
15
+ from torch.utils.checkpoint import checkpoint
16
+ import math
17
+
18
+ try:
19
+ from typing import Optional, NamedTuple, List, Protocol
20
+ except ImportError:
21
+ from typing import Optional, NamedTuple, List
22
+ from typing_extensions import Protocol
23
+
24
+ from torch import Tensor
25
+ from typing import List
26
+
27
+ from comfy import model_management
28
+
29
+ def dynamic_slice(
30
+ x: Tensor,
31
+ starts: List[int],
32
+ sizes: List[int],
33
+ ) -> Tensor:
34
+ slicing = [slice(start, start + size) for start, size in zip(starts, sizes)]
35
+ return x[slicing]
36
+
37
+ class AttnChunk(NamedTuple):
38
+ exp_values: Tensor
39
+ exp_weights_sum: Tensor
40
+ max_score: Tensor
41
+
42
+ class SummarizeChunk(Protocol):
43
+ @staticmethod
44
+ def __call__(
45
+ query: Tensor,
46
+ key_t: Tensor,
47
+ value: Tensor,
48
+ ) -> AttnChunk: ...
49
+
50
+ class ComputeQueryChunkAttn(Protocol):
51
+ @staticmethod
52
+ def __call__(
53
+ query: Tensor,
54
+ key_t: Tensor,
55
+ value: Tensor,
56
+ ) -> Tensor: ...
57
+
58
+ def _summarize_chunk(
59
+ query: Tensor,
60
+ key_t: Tensor,
61
+ value: Tensor,
62
+ scale: float,
63
+ upcast_attention: bool,
64
+ ) -> AttnChunk:
65
+ if upcast_attention:
66
+ with torch.autocast(enabled=False, device_type = 'cuda'):
67
+ query = query.float()
68
+ key_t = key_t.float()
69
+ attn_weights = torch.baddbmm(
70
+ torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
71
+ query,
72
+ key_t,
73
+ alpha=scale,
74
+ beta=0,
75
+ )
76
+ else:
77
+ attn_weights = torch.baddbmm(
78
+ torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
79
+ query,
80
+ key_t,
81
+ alpha=scale,
82
+ beta=0,
83
+ )
84
+ max_score, _ = torch.max(attn_weights, -1, keepdim=True)
85
+ max_score = max_score.detach()
86
+ torch.exp(attn_weights - max_score, out=attn_weights)
87
+ exp_weights = attn_weights.to(value.dtype)
88
+ exp_values = torch.bmm(exp_weights, value)
89
+ max_score = max_score.squeeze(-1)
90
+ return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
91
+
92
+ def _query_chunk_attention(
93
+ query: Tensor,
94
+ key_t: Tensor,
95
+ value: Tensor,
96
+ summarize_chunk: SummarizeChunk,
97
+ kv_chunk_size: int,
98
+ ) -> Tensor:
99
+ batch_x_heads, k_channels_per_head, k_tokens = key_t.shape
100
+ _, _, v_channels_per_head = value.shape
101
+
102
+ def chunk_scanner(chunk_idx: int) -> AttnChunk:
103
+ key_chunk = dynamic_slice(
104
+ key_t,
105
+ (0, 0, chunk_idx),
106
+ (batch_x_heads, k_channels_per_head, kv_chunk_size)
107
+ )
108
+ value_chunk = dynamic_slice(
109
+ value,
110
+ (0, chunk_idx, 0),
111
+ (batch_x_heads, kv_chunk_size, v_channels_per_head)
112
+ )
113
+ return summarize_chunk(query, key_chunk, value_chunk)
114
+
115
+ chunks: List[AttnChunk] = [
116
+ chunk_scanner(chunk) for chunk in torch.arange(0, k_tokens, kv_chunk_size)
117
+ ]
118
+ acc_chunk = AttnChunk(*map(torch.stack, zip(*chunks)))
119
+ chunk_values, chunk_weights, chunk_max = acc_chunk
120
+
121
+ global_max, _ = torch.max(chunk_max, 0, keepdim=True)
122
+ max_diffs = torch.exp(chunk_max - global_max)
123
+ chunk_values *= torch.unsqueeze(max_diffs, -1)
124
+ chunk_weights *= max_diffs
125
+
126
+ all_values = chunk_values.sum(dim=0)
127
+ all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0)
128
+ return all_values / all_weights
129
+
130
+ # TODO: refactor CrossAttention#get_attention_scores to share code with this
131
+ def _get_attention_scores_no_kv_chunking(
132
+ query: Tensor,
133
+ key_t: Tensor,
134
+ value: Tensor,
135
+ scale: float,
136
+ upcast_attention: bool,
137
+ ) -> Tensor:
138
+ if upcast_attention:
139
+ with torch.autocast(enabled=False, device_type = 'cuda'):
140
+ query = query.float()
141
+ key_t = key_t.float()
142
+ attn_scores = torch.baddbmm(
143
+ torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
144
+ query,
145
+ key_t,
146
+ alpha=scale,
147
+ beta=0,
148
+ )
149
+ else:
150
+ attn_scores = torch.baddbmm(
151
+ torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
152
+ query,
153
+ key_t,
154
+ alpha=scale,
155
+ beta=0,
156
+ )
157
+
158
+ try:
159
+ attn_probs = attn_scores.softmax(dim=-1)
160
+ del attn_scores
161
+ except model_management.OOM_EXCEPTION:
162
+ print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
163
+ attn_scores -= attn_scores.max(dim=-1, keepdim=True).values
164
+ torch.exp(attn_scores, out=attn_scores)
165
+ summed = torch.sum(attn_scores, dim=-1, keepdim=True)
166
+ attn_scores /= summed
167
+ attn_probs = attn_scores
168
+
169
+ hidden_states_slice = torch.bmm(attn_probs.to(value.dtype), value)
170
+ return hidden_states_slice
171
+
172
+ class ScannedChunk(NamedTuple):
173
+ chunk_idx: int
174
+ attn_chunk: AttnChunk
175
+
176
+ def efficient_dot_product_attention(
177
+ query: Tensor,
178
+ key_t: Tensor,
179
+ value: Tensor,
180
+ query_chunk_size=1024,
181
+ kv_chunk_size: Optional[int] = None,
182
+ kv_chunk_size_min: Optional[int] = None,
183
+ use_checkpoint=True,
184
+ upcast_attention=False,
185
+ ):
186
+ """Computes efficient dot-product attention given query, transposed key, and value.
187
+ This is efficient version of attention presented in
188
+ https://arxiv.org/abs/2112.05682v2 which comes with O(sqrt(n)) memory requirements.
189
+ Args:
190
+ query: queries for calculating attention with shape of
191
+ `[batch * num_heads, tokens, channels_per_head]`.
192
+ key_t: keys for calculating attention with shape of
193
+ `[batch * num_heads, channels_per_head, tokens]`.
194
+ value: values to be used in attention with shape of
195
+ `[batch * num_heads, tokens, channels_per_head]`.
196
+ query_chunk_size: int: query chunks size
197
+ kv_chunk_size: Optional[int]: key/value chunks size. if None: defaults to sqrt(key_tokens)
198
+ kv_chunk_size_min: Optional[int]: key/value minimum chunk size. only considered when kv_chunk_size is None. changes `sqrt(key_tokens)` into `max(sqrt(key_tokens), kv_chunk_size_min)`, to ensure our chunk sizes don't get too small (smaller chunks = more chunks = less concurrent work done).
199
+ use_checkpoint: bool: whether to use checkpointing (recommended True for training, False for inference)
200
+ Returns:
201
+ Output of shape `[batch * num_heads, query_tokens, channels_per_head]`.
202
+ """
203
+ batch_x_heads, q_tokens, q_channels_per_head = query.shape
204
+ _, _, k_tokens = key_t.shape
205
+ scale = q_channels_per_head ** -0.5
206
+
207
+ kv_chunk_size = min(kv_chunk_size or int(math.sqrt(k_tokens)), k_tokens)
208
+ if kv_chunk_size_min is not None:
209
+ kv_chunk_size = max(kv_chunk_size, kv_chunk_size_min)
210
+
211
+ def get_query_chunk(chunk_idx: int) -> Tensor:
212
+ return dynamic_slice(
213
+ query,
214
+ (0, chunk_idx, 0),
215
+ (batch_x_heads, min(query_chunk_size, q_tokens), q_channels_per_head)
216
+ )
217
+
218
+ summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale, upcast_attention=upcast_attention)
219
+ summarize_chunk: SummarizeChunk = partial(checkpoint, summarize_chunk) if use_checkpoint else summarize_chunk
220
+ compute_query_chunk_attn: ComputeQueryChunkAttn = partial(
221
+ _get_attention_scores_no_kv_chunking,
222
+ scale=scale,
223
+ upcast_attention=upcast_attention
224
+ ) if k_tokens <= kv_chunk_size else (
225
+ # fast-path for when there's just 1 key-value chunk per query chunk (this is just sliced attention btw)
226
+ partial(
227
+ _query_chunk_attention,
228
+ kv_chunk_size=kv_chunk_size,
229
+ summarize_chunk=summarize_chunk,
230
+ )
231
+ )
232
+
233
+ if q_tokens <= query_chunk_size:
234
+ # fast-path for when there's just 1 query chunk
235
+ return compute_query_chunk_attn(
236
+ query=query,
237
+ key_t=key_t,
238
+ value=value,
239
+ )
240
+
241
+ # TODO: maybe we should use torch.empty_like(query) to allocate storage in-advance,
242
+ # and pass slices to be mutated, instead of torch.cat()ing the returned slices
243
+ res = torch.cat([
244
+ compute_query_chunk_attn(
245
+ query=get_query_chunk(i * query_chunk_size),
246
+ key_t=key_t,
247
+ value=value,
248
+ ) for i in range(math.ceil(q_tokens / query_chunk_size))
249
+ ], dim=1)
250
+ return res
comfy/ldm/util.py ADDED
@@ -0,0 +1,197 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import importlib
2
+
3
+ import torch
4
+ from torch import optim
5
+ import numpy as np
6
+
7
+ from inspect import isfunction
8
+ from PIL import Image, ImageDraw, ImageFont
9
+
10
+
11
+ def log_txt_as_img(wh, xc, size=10):
12
+ # wh a tuple of (width, height)
13
+ # xc a list of captions to plot
14
+ b = len(xc)
15
+ txts = list()
16
+ for bi in range(b):
17
+ txt = Image.new("RGB", wh, color="white")
18
+ draw = ImageDraw.Draw(txt)
19
+ font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
20
+ nc = int(40 * (wh[0] / 256))
21
+ lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
22
+
23
+ try:
24
+ draw.text((0, 0), lines, fill="black", font=font)
25
+ except UnicodeEncodeError:
26
+ print("Cant encode string for logging. Skipping.")
27
+
28
+ txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
29
+ txts.append(txt)
30
+ txts = np.stack(txts)
31
+ txts = torch.tensor(txts)
32
+ return txts
33
+
34
+
35
+ def ismap(x):
36
+ if not isinstance(x, torch.Tensor):
37
+ return False
38
+ return (len(x.shape) == 4) and (x.shape[1] > 3)
39
+
40
+
41
+ def isimage(x):
42
+ if not isinstance(x,torch.Tensor):
43
+ return False
44
+ return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
45
+
46
+
47
+ def exists(x):
48
+ return x is not None
49
+
50
+
51
+ def default(val, d):
52
+ if exists(val):
53
+ return val
54
+ return d() if isfunction(d) else d
55
+
56
+
57
+ def mean_flat(tensor):
58
+ """
59
+ https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
60
+ Take the mean over all non-batch dimensions.
61
+ """
62
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
63
+
64
+
65
+ def count_params(model, verbose=False):
66
+ total_params = sum(p.numel() for p in model.parameters())
67
+ if verbose:
68
+ print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
69
+ return total_params
70
+
71
+
72
+ def instantiate_from_config(config):
73
+ if not "target" in config:
74
+ if config == '__is_first_stage__':
75
+ return None
76
+ elif config == "__is_unconditional__":
77
+ return None
78
+ raise KeyError("Expected key `target` to instantiate.")
79
+ return get_obj_from_str(config["target"])(**config.get("params", dict()))
80
+
81
+
82
+ def get_obj_from_str(string, reload=False):
83
+ module, cls = string.rsplit(".", 1)
84
+ if reload:
85
+ module_imp = importlib.import_module(module)
86
+ importlib.reload(module_imp)
87
+ return getattr(importlib.import_module(module, package=None), cls)
88
+
89
+
90
+ class AdamWwithEMAandWings(optim.Optimizer):
91
+ # credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
92
+ def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, # TODO: check hyperparameters before using
93
+ weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999, # ema decay to match previous code
94
+ ema_power=1., param_names=()):
95
+ """AdamW that saves EMA versions of the parameters."""
96
+ if not 0.0 <= lr:
97
+ raise ValueError("Invalid learning rate: {}".format(lr))
98
+ if not 0.0 <= eps:
99
+ raise ValueError("Invalid epsilon value: {}".format(eps))
100
+ if not 0.0 <= betas[0] < 1.0:
101
+ raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
102
+ if not 0.0 <= betas[1] < 1.0:
103
+ raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
104
+ if not 0.0 <= weight_decay:
105
+ raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
106
+ if not 0.0 <= ema_decay <= 1.0:
107
+ raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
108
+ defaults = dict(lr=lr, betas=betas, eps=eps,
109
+ weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
110
+ ema_power=ema_power, param_names=param_names)
111
+ super().__init__(params, defaults)
112
+
113
+ def __setstate__(self, state):
114
+ super().__setstate__(state)
115
+ for group in self.param_groups:
116
+ group.setdefault('amsgrad', False)
117
+
118
+ @torch.no_grad()
119
+ def step(self, closure=None):
120
+ """Performs a single optimization step.
121
+ Args:
122
+ closure (callable, optional): A closure that reevaluates the model
123
+ and returns the loss.
124
+ """
125
+ loss = None
126
+ if closure is not None:
127
+ with torch.enable_grad():
128
+ loss = closure()
129
+
130
+ for group in self.param_groups:
131
+ params_with_grad = []
132
+ grads = []
133
+ exp_avgs = []
134
+ exp_avg_sqs = []
135
+ ema_params_with_grad = []
136
+ state_sums = []
137
+ max_exp_avg_sqs = []
138
+ state_steps = []
139
+ amsgrad = group['amsgrad']
140
+ beta1, beta2 = group['betas']
141
+ ema_decay = group['ema_decay']
142
+ ema_power = group['ema_power']
143
+
144
+ for p in group['params']:
145
+ if p.grad is None:
146
+ continue
147
+ params_with_grad.append(p)
148
+ if p.grad.is_sparse:
149
+ raise RuntimeError('AdamW does not support sparse gradients')
150
+ grads.append(p.grad)
151
+
152
+ state = self.state[p]
153
+
154
+ # State initialization
155
+ if len(state) == 0:
156
+ state['step'] = 0
157
+ # Exponential moving average of gradient values
158
+ state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
159
+ # Exponential moving average of squared gradient values
160
+ state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
161
+ if amsgrad:
162
+ # Maintains max of all exp. moving avg. of sq. grad. values
163
+ state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
164
+ # Exponential moving average of parameter values
165
+ state['param_exp_avg'] = p.detach().float().clone()
166
+
167
+ exp_avgs.append(state['exp_avg'])
168
+ exp_avg_sqs.append(state['exp_avg_sq'])
169
+ ema_params_with_grad.append(state['param_exp_avg'])
170
+
171
+ if amsgrad:
172
+ max_exp_avg_sqs.append(state['max_exp_avg_sq'])
173
+
174
+ # update the steps for each param group update
175
+ state['step'] += 1
176
+ # record the step after step update
177
+ state_steps.append(state['step'])
178
+
179
+ optim._functional.adamw(params_with_grad,
180
+ grads,
181
+ exp_avgs,
182
+ exp_avg_sqs,
183
+ max_exp_avg_sqs,
184
+ state_steps,
185
+ amsgrad=amsgrad,
186
+ beta1=beta1,
187
+ beta2=beta2,
188
+ lr=group['lr'],
189
+ weight_decay=group['weight_decay'],
190
+ eps=group['eps'],
191
+ maximize=False)
192
+
193
+ cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
194
+ for param, ema_param in zip(params_with_grad, ema_params_with_grad):
195
+ ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
196
+
197
+ return loss
comfy/lora.py ADDED
@@ -0,0 +1,199 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import comfy.utils
2
+
3
+ LORA_CLIP_MAP = {
4
+ "mlp.fc1": "mlp_fc1",
5
+ "mlp.fc2": "mlp_fc2",
6
+ "self_attn.k_proj": "self_attn_k_proj",
7
+ "self_attn.q_proj": "self_attn_q_proj",
8
+ "self_attn.v_proj": "self_attn_v_proj",
9
+ "self_attn.out_proj": "self_attn_out_proj",
10
+ }
11
+
12
+
13
+ def load_lora(lora, to_load):
14
+ patch_dict = {}
15
+ loaded_keys = set()
16
+ for x in to_load:
17
+ alpha_name = "{}.alpha".format(x)
18
+ alpha = None
19
+ if alpha_name in lora.keys():
20
+ alpha = lora[alpha_name].item()
21
+ loaded_keys.add(alpha_name)
22
+
23
+ regular_lora = "{}.lora_up.weight".format(x)
24
+ diffusers_lora = "{}_lora.up.weight".format(x)
25
+ transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
26
+ A_name = None
27
+
28
+ if regular_lora in lora.keys():
29
+ A_name = regular_lora
30
+ B_name = "{}.lora_down.weight".format(x)
31
+ mid_name = "{}.lora_mid.weight".format(x)
32
+ elif diffusers_lora in lora.keys():
33
+ A_name = diffusers_lora
34
+ B_name = "{}_lora.down.weight".format(x)
35
+ mid_name = None
36
+ elif transformers_lora in lora.keys():
37
+ A_name = transformers_lora
38
+ B_name ="{}.lora_linear_layer.down.weight".format(x)
39
+ mid_name = None
40
+
41
+ if A_name is not None:
42
+ mid = None
43
+ if mid_name is not None and mid_name in lora.keys():
44
+ mid = lora[mid_name]
45
+ loaded_keys.add(mid_name)
46
+ patch_dict[to_load[x]] = (lora[A_name], lora[B_name], alpha, mid)
47
+ loaded_keys.add(A_name)
48
+ loaded_keys.add(B_name)
49
+
50
+
51
+ ######## loha
52
+ hada_w1_a_name = "{}.hada_w1_a".format(x)
53
+ hada_w1_b_name = "{}.hada_w1_b".format(x)
54
+ hada_w2_a_name = "{}.hada_w2_a".format(x)
55
+ hada_w2_b_name = "{}.hada_w2_b".format(x)
56
+ hada_t1_name = "{}.hada_t1".format(x)
57
+ hada_t2_name = "{}.hada_t2".format(x)
58
+ if hada_w1_a_name in lora.keys():
59
+ hada_t1 = None
60
+ hada_t2 = None
61
+ if hada_t1_name in lora.keys():
62
+ hada_t1 = lora[hada_t1_name]
63
+ hada_t2 = lora[hada_t2_name]
64
+ loaded_keys.add(hada_t1_name)
65
+ loaded_keys.add(hada_t2_name)
66
+
67
+ patch_dict[to_load[x]] = (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2)
68
+ loaded_keys.add(hada_w1_a_name)
69
+ loaded_keys.add(hada_w1_b_name)
70
+ loaded_keys.add(hada_w2_a_name)
71
+ loaded_keys.add(hada_w2_b_name)
72
+
73
+
74
+ ######## lokr
75
+ lokr_w1_name = "{}.lokr_w1".format(x)
76
+ lokr_w2_name = "{}.lokr_w2".format(x)
77
+ lokr_w1_a_name = "{}.lokr_w1_a".format(x)
78
+ lokr_w1_b_name = "{}.lokr_w1_b".format(x)
79
+ lokr_t2_name = "{}.lokr_t2".format(x)
80
+ lokr_w2_a_name = "{}.lokr_w2_a".format(x)
81
+ lokr_w2_b_name = "{}.lokr_w2_b".format(x)
82
+
83
+ lokr_w1 = None
84
+ if lokr_w1_name in lora.keys():
85
+ lokr_w1 = lora[lokr_w1_name]
86
+ loaded_keys.add(lokr_w1_name)
87
+
88
+ lokr_w2 = None
89
+ if lokr_w2_name in lora.keys():
90
+ lokr_w2 = lora[lokr_w2_name]
91
+ loaded_keys.add(lokr_w2_name)
92
+
93
+ lokr_w1_a = None
94
+ if lokr_w1_a_name in lora.keys():
95
+ lokr_w1_a = lora[lokr_w1_a_name]
96
+ loaded_keys.add(lokr_w1_a_name)
97
+
98
+ lokr_w1_b = None
99
+ if lokr_w1_b_name in lora.keys():
100
+ lokr_w1_b = lora[lokr_w1_b_name]
101
+ loaded_keys.add(lokr_w1_b_name)
102
+
103
+ lokr_w2_a = None
104
+ if lokr_w2_a_name in lora.keys():
105
+ lokr_w2_a = lora[lokr_w2_a_name]
106
+ loaded_keys.add(lokr_w2_a_name)
107
+
108
+ lokr_w2_b = None
109
+ if lokr_w2_b_name in lora.keys():
110
+ lokr_w2_b = lora[lokr_w2_b_name]
111
+ loaded_keys.add(lokr_w2_b_name)
112
+
113
+ lokr_t2 = None
114
+ if lokr_t2_name in lora.keys():
115
+ lokr_t2 = lora[lokr_t2_name]
116
+ loaded_keys.add(lokr_t2_name)
117
+
118
+ if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
119
+ patch_dict[to_load[x]] = (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2)
120
+
121
+
122
+ w_norm_name = "{}.w_norm".format(x)
123
+ b_norm_name = "{}.b_norm".format(x)
124
+ w_norm = lora.get(w_norm_name, None)
125
+ b_norm = lora.get(b_norm_name, None)
126
+
127
+ if w_norm is not None:
128
+ loaded_keys.add(w_norm_name)
129
+ patch_dict[to_load[x]] = (w_norm,)
130
+ if b_norm is not None:
131
+ loaded_keys.add(b_norm_name)
132
+ patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = (b_norm,)
133
+
134
+ for x in lora.keys():
135
+ if x not in loaded_keys:
136
+ print("lora key not loaded", x)
137
+ return patch_dict
138
+
139
+ def model_lora_keys_clip(model, key_map={}):
140
+ sdk = model.state_dict().keys()
141
+
142
+ text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}"
143
+ clip_l_present = False
144
+ for b in range(32):
145
+ for c in LORA_CLIP_MAP:
146
+ k = "transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
147
+ if k in sdk:
148
+ lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])
149
+ key_map[lora_key] = k
150
+ lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c])
151
+ key_map[lora_key] = k
152
+ lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
153
+ key_map[lora_key] = k
154
+
155
+ k = "clip_l.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
156
+ if k in sdk:
157
+ lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
158
+ key_map[lora_key] = k
159
+ clip_l_present = True
160
+ lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
161
+ key_map[lora_key] = k
162
+
163
+ k = "clip_g.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
164
+ if k in sdk:
165
+ if clip_l_present:
166
+ lora_key = "lora_te2_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
167
+ key_map[lora_key] = k
168
+ lora_key = "text_encoder_2.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
169
+ key_map[lora_key] = k
170
+ else:
171
+ lora_key = "lora_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #TODO: test if this is correct for SDXL-Refiner
172
+ key_map[lora_key] = k
173
+ lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
174
+ key_map[lora_key] = k
175
+
176
+ return key_map
177
+
178
+ def model_lora_keys_unet(model, key_map={}):
179
+ sdk = model.state_dict().keys()
180
+
181
+ for k in sdk:
182
+ if k.startswith("diffusion_model.") and k.endswith(".weight"):
183
+ key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
184
+ key_map["lora_unet_{}".format(key_lora)] = k
185
+
186
+ diffusers_keys = comfy.utils.unet_to_diffusers(model.model_config.unet_config)
187
+ for k in diffusers_keys:
188
+ if k.endswith(".weight"):
189
+ unet_key = "diffusion_model.{}".format(diffusers_keys[k])
190
+ key_lora = k[:-len(".weight")].replace(".", "_")
191
+ key_map["lora_unet_{}".format(key_lora)] = unet_key
192
+
193
+ diffusers_lora_prefix = ["", "unet."]
194
+ for p in diffusers_lora_prefix:
195
+ diffusers_lora_key = "{}{}".format(p, k[:-len(".weight")].replace(".to_", ".processor.to_"))
196
+ if diffusers_lora_key.endswith(".to_out.0"):
197
+ diffusers_lora_key = diffusers_lora_key[:-2]
198
+ key_map[diffusers_lora_key] = unet_key
199
+ return key_map
comfy/model_base.py ADDED
@@ -0,0 +1,210 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel
3
+ from comfy.ldm.modules.encoders.noise_aug_modules import CLIPEmbeddingNoiseAugmentation
4
+ from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule
5
+ from comfy.ldm.modules.diffusionmodules.openaimodel import Timestep
6
+ import comfy.model_management
7
+ import numpy as np
8
+ from enum import Enum
9
+ from . import utils
10
+
11
+ class ModelType(Enum):
12
+ EPS = 1
13
+ V_PREDICTION = 2
14
+
15
+ class BaseModel(torch.nn.Module):
16
+ def __init__(self, model_config, model_type=ModelType.EPS, device=None):
17
+ super().__init__()
18
+
19
+ unet_config = model_config.unet_config
20
+ self.latent_format = model_config.latent_format
21
+ self.model_config = model_config
22
+ self.register_schedule(given_betas=None, beta_schedule=model_config.beta_schedule, timesteps=1000, linear_start=0.00085, linear_end=0.012, cosine_s=8e-3)
23
+ if not unet_config.get("disable_unet_model_creation", False):
24
+ self.diffusion_model = UNetModel(**unet_config, device=device)
25
+ self.model_type = model_type
26
+ self.adm_channels = unet_config.get("adm_in_channels", None)
27
+ if self.adm_channels is None:
28
+ self.adm_channels = 0
29
+ print("model_type", model_type.name)
30
+ print("adm", self.adm_channels)
31
+
32
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
33
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
34
+ if given_betas is not None:
35
+ betas = given_betas
36
+ else:
37
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
38
+ alphas = 1. - betas
39
+ alphas_cumprod = np.cumprod(alphas, axis=0)
40
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
41
+
42
+ timesteps, = betas.shape
43
+ self.num_timesteps = int(timesteps)
44
+ self.linear_start = linear_start
45
+ self.linear_end = linear_end
46
+
47
+ self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32))
48
+ self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32))
49
+ self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))
50
+
51
+ def apply_model(self, x, t, c_concat=None, c_crossattn=None, c_adm=None, control=None, transformer_options={}):
52
+ if c_concat is not None:
53
+ xc = torch.cat([x] + [c_concat], dim=1)
54
+ else:
55
+ xc = x
56
+ context = c_crossattn
57
+ dtype = self.get_dtype()
58
+ xc = xc.to(dtype)
59
+ t = t.to(dtype)
60
+ context = context.to(dtype)
61
+ if c_adm is not None:
62
+ c_adm = c_adm.to(dtype)
63
+ return self.diffusion_model(xc, t, context=context, y=c_adm, control=control, transformer_options=transformer_options).float()
64
+
65
+ def get_dtype(self):
66
+ return self.diffusion_model.dtype
67
+
68
+ def is_adm(self):
69
+ return self.adm_channels > 0
70
+
71
+ def encode_adm(self, **kwargs):
72
+ return None
73
+
74
+ def load_model_weights(self, sd, unet_prefix=""):
75
+ to_load = {}
76
+ keys = list(sd.keys())
77
+ for k in keys:
78
+ if k.startswith(unet_prefix):
79
+ to_load[k[len(unet_prefix):]] = sd.pop(k)
80
+
81
+ m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
82
+ if len(m) > 0:
83
+ print("unet missing:", m)
84
+
85
+ if len(u) > 0:
86
+ print("unet unexpected:", u)
87
+ del to_load
88
+ return self
89
+
90
+ def process_latent_in(self, latent):
91
+ return self.latent_format.process_in(latent)
92
+
93
+ def process_latent_out(self, latent):
94
+ return self.latent_format.process_out(latent)
95
+
96
+ def state_dict_for_saving(self, clip_state_dict, vae_state_dict):
97
+ clip_state_dict = self.model_config.process_clip_state_dict_for_saving(clip_state_dict)
98
+ unet_sd = self.diffusion_model.state_dict()
99
+ unet_state_dict = {}
100
+ for k in unet_sd:
101
+ unet_state_dict[k] = comfy.model_management.resolve_lowvram_weight(unet_sd[k], self.diffusion_model, k)
102
+
103
+ unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
104
+ vae_state_dict = self.model_config.process_vae_state_dict_for_saving(vae_state_dict)
105
+ if self.get_dtype() == torch.float16:
106
+ clip_state_dict = utils.convert_sd_to(clip_state_dict, torch.float16)
107
+ vae_state_dict = utils.convert_sd_to(vae_state_dict, torch.float16)
108
+
109
+ if self.model_type == ModelType.V_PREDICTION:
110
+ unet_state_dict["v_pred"] = torch.tensor([])
111
+
112
+ return {**unet_state_dict, **vae_state_dict, **clip_state_dict}
113
+
114
+ def set_inpaint(self):
115
+ self.concat_keys = ("mask", "masked_image")
116
+
117
+ def unclip_adm(unclip_conditioning, device, noise_augmentor, noise_augment_merge=0.0):
118
+ adm_inputs = []
119
+ weights = []
120
+ noise_aug = []
121
+ for unclip_cond in unclip_conditioning:
122
+ for adm_cond in unclip_cond["clip_vision_output"].image_embeds:
123
+ weight = unclip_cond["strength"]
124
+ noise_augment = unclip_cond["noise_augmentation"]
125
+ noise_level = round((noise_augmentor.max_noise_level - 1) * noise_augment)
126
+ c_adm, noise_level_emb = noise_augmentor(adm_cond.to(device), noise_level=torch.tensor([noise_level], device=device))
127
+ adm_out = torch.cat((c_adm, noise_level_emb), 1) * weight
128
+ weights.append(weight)
129
+ noise_aug.append(noise_augment)
130
+ adm_inputs.append(adm_out)
131
+
132
+ if len(noise_aug) > 1:
133
+ adm_out = torch.stack(adm_inputs).sum(0)
134
+ noise_augment = noise_augment_merge
135
+ noise_level = round((noise_augmentor.max_noise_level - 1) * noise_augment)
136
+ c_adm, noise_level_emb = noise_augmentor(adm_out[:, :noise_augmentor.time_embed.dim], noise_level=torch.tensor([noise_level], device=device))
137
+ adm_out = torch.cat((c_adm, noise_level_emb), 1)
138
+
139
+ return adm_out
140
+
141
+ class SD21UNCLIP(BaseModel):
142
+ def __init__(self, model_config, noise_aug_config, model_type=ModelType.V_PREDICTION, device=None):
143
+ super().__init__(model_config, model_type, device=device)
144
+ self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**noise_aug_config)
145
+
146
+ def encode_adm(self, **kwargs):
147
+ unclip_conditioning = kwargs.get("unclip_conditioning", None)
148
+ device = kwargs["device"]
149
+ if unclip_conditioning is None:
150
+ return torch.zeros((1, self.adm_channels))
151
+ else:
152
+ return unclip_adm(unclip_conditioning, device, self.noise_augmentor, kwargs.get("unclip_noise_augment_merge", 0.05))
153
+
154
+ def sdxl_pooled(args, noise_augmentor):
155
+ if "unclip_conditioning" in args:
156
+ return unclip_adm(args.get("unclip_conditioning", None), args["device"], noise_augmentor)[:,:1280]
157
+ else:
158
+ return args["pooled_output"]
159
+
160
+ class SDXLRefiner(BaseModel):
161
+ def __init__(self, model_config, model_type=ModelType.EPS, device=None):
162
+ super().__init__(model_config, model_type, device=device)
163
+ self.embedder = Timestep(256)
164
+ self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**{"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1280})
165
+
166
+ def encode_adm(self, **kwargs):
167
+ clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
168
+ width = kwargs.get("width", 768)
169
+ height = kwargs.get("height", 768)
170
+ crop_w = kwargs.get("crop_w", 0)
171
+ crop_h = kwargs.get("crop_h", 0)
172
+
173
+ if kwargs.get("prompt_type", "") == "negative":
174
+ aesthetic_score = kwargs.get("aesthetic_score", 2.5)
175
+ else:
176
+ aesthetic_score = kwargs.get("aesthetic_score", 6)
177
+
178
+ out = []
179
+ out.append(self.embedder(torch.Tensor([height])))
180
+ out.append(self.embedder(torch.Tensor([width])))
181
+ out.append(self.embedder(torch.Tensor([crop_h])))
182
+ out.append(self.embedder(torch.Tensor([crop_w])))
183
+ out.append(self.embedder(torch.Tensor([aesthetic_score])))
184
+ flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
185
+ return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
186
+
187
+ class SDXL(BaseModel):
188
+ def __init__(self, model_config, model_type=ModelType.EPS, device=None):
189
+ super().__init__(model_config, model_type, device=device)
190
+ self.embedder = Timestep(256)
191
+ self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**{"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1280})
192
+
193
+ def encode_adm(self, **kwargs):
194
+ clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
195
+ width = kwargs.get("width", 768)
196
+ height = kwargs.get("height", 768)
197
+ crop_w = kwargs.get("crop_w", 0)
198
+ crop_h = kwargs.get("crop_h", 0)
199
+ target_width = kwargs.get("target_width", width)
200
+ target_height = kwargs.get("target_height", height)
201
+
202
+ out = []
203
+ out.append(self.embedder(torch.Tensor([height])))
204
+ out.append(self.embedder(torch.Tensor([width])))
205
+ out.append(self.embedder(torch.Tensor([crop_h])))
206
+ out.append(self.embedder(torch.Tensor([crop_w])))
207
+ out.append(self.embedder(torch.Tensor([target_height])))
208
+ out.append(self.embedder(torch.Tensor([target_width])))
209
+ flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
210
+ return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
comfy/model_detection.py ADDED
@@ -0,0 +1,210 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import comfy.supported_models
2
+ import comfy.supported_models_base
3
+
4
+ def count_blocks(state_dict_keys, prefix_string):
5
+ count = 0
6
+ while True:
7
+ c = False
8
+ for k in state_dict_keys:
9
+ if k.startswith(prefix_string.format(count)):
10
+ c = True
11
+ break
12
+ if c == False:
13
+ break
14
+ count += 1
15
+ return count
16
+
17
+ def detect_unet_config(state_dict, key_prefix, use_fp16):
18
+ state_dict_keys = list(state_dict.keys())
19
+
20
+ unet_config = {
21
+ "use_checkpoint": False,
22
+ "image_size": 32,
23
+ "out_channels": 4,
24
+ "use_spatial_transformer": True,
25
+ "legacy": False
26
+ }
27
+
28
+ y_input = '{}label_emb.0.0.weight'.format(key_prefix)
29
+ if y_input in state_dict_keys:
30
+ unet_config["num_classes"] = "sequential"
31
+ unet_config["adm_in_channels"] = state_dict[y_input].shape[1]
32
+ else:
33
+ unet_config["adm_in_channels"] = None
34
+
35
+ unet_config["use_fp16"] = use_fp16
36
+ model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0]
37
+ in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1]
38
+
39
+ num_res_blocks = []
40
+ channel_mult = []
41
+ attention_resolutions = []
42
+ transformer_depth = []
43
+ context_dim = None
44
+ use_linear_in_transformer = False
45
+
46
+
47
+ current_res = 1
48
+ count = 0
49
+
50
+ last_res_blocks = 0
51
+ last_transformer_depth = 0
52
+ last_channel_mult = 0
53
+
54
+ while True:
55
+ prefix = '{}input_blocks.{}.'.format(key_prefix, count)
56
+ block_keys = sorted(list(filter(lambda a: a.startswith(prefix), state_dict_keys)))
57
+ if len(block_keys) == 0:
58
+ break
59
+
60
+ if "{}0.op.weight".format(prefix) in block_keys: #new layer
61
+ if last_transformer_depth > 0:
62
+ attention_resolutions.append(current_res)
63
+ transformer_depth.append(last_transformer_depth)
64
+ num_res_blocks.append(last_res_blocks)
65
+ channel_mult.append(last_channel_mult)
66
+
67
+ current_res *= 2
68
+ last_res_blocks = 0
69
+ last_transformer_depth = 0
70
+ last_channel_mult = 0
71
+ else:
72
+ res_block_prefix = "{}0.in_layers.0.weight".format(prefix)
73
+ if res_block_prefix in block_keys:
74
+ last_res_blocks += 1
75
+ last_channel_mult = state_dict["{}0.out_layers.3.weight".format(prefix)].shape[0] // model_channels
76
+
77
+ transformer_prefix = prefix + "1.transformer_blocks."
78
+ transformer_keys = sorted(list(filter(lambda a: a.startswith(transformer_prefix), state_dict_keys)))
79
+ if len(transformer_keys) > 0:
80
+ last_transformer_depth = count_blocks(state_dict_keys, transformer_prefix + '{}')
81
+ if context_dim is None:
82
+ context_dim = state_dict['{}0.attn2.to_k.weight'.format(transformer_prefix)].shape[1]
83
+ use_linear_in_transformer = len(state_dict['{}1.proj_in.weight'.format(prefix)].shape) == 2
84
+
85
+ count += 1
86
+
87
+ if last_transformer_depth > 0:
88
+ attention_resolutions.append(current_res)
89
+ transformer_depth.append(last_transformer_depth)
90
+ num_res_blocks.append(last_res_blocks)
91
+ channel_mult.append(last_channel_mult)
92
+ transformer_depth_middle = count_blocks(state_dict_keys, '{}middle_block.1.transformer_blocks.'.format(key_prefix) + '{}')
93
+
94
+ if len(set(num_res_blocks)) == 1:
95
+ num_res_blocks = num_res_blocks[0]
96
+
97
+ if len(set(transformer_depth)) == 1:
98
+ transformer_depth = transformer_depth[0]
99
+
100
+ unet_config["in_channels"] = in_channels
101
+ unet_config["model_channels"] = model_channels
102
+ unet_config["num_res_blocks"] = num_res_blocks
103
+ unet_config["attention_resolutions"] = attention_resolutions
104
+ unet_config["transformer_depth"] = transformer_depth
105
+ unet_config["channel_mult"] = channel_mult
106
+ unet_config["transformer_depth_middle"] = transformer_depth_middle
107
+ unet_config['use_linear_in_transformer'] = use_linear_in_transformer
108
+ unet_config["context_dim"] = context_dim
109
+ return unet_config
110
+
111
+ def model_config_from_unet_config(unet_config):
112
+ for model_config in comfy.supported_models.models:
113
+ if model_config.matches(unet_config):
114
+ return model_config(unet_config)
115
+
116
+ print("no match", unet_config)
117
+ return None
118
+
119
+ def model_config_from_unet(state_dict, unet_key_prefix, use_fp16, use_base_if_no_match=False):
120
+ unet_config = detect_unet_config(state_dict, unet_key_prefix, use_fp16)
121
+ model_config = model_config_from_unet_config(unet_config)
122
+ if model_config is None and use_base_if_no_match:
123
+ return comfy.supported_models_base.BASE(unet_config)
124
+ else:
125
+ return model_config
126
+
127
+ def unet_config_from_diffusers_unet(state_dict, use_fp16):
128
+ match = {}
129
+ attention_resolutions = []
130
+
131
+ attn_res = 1
132
+ for i in range(5):
133
+ k = "down_blocks.{}.attentions.1.transformer_blocks.0.attn2.to_k.weight".format(i)
134
+ if k in state_dict:
135
+ match["context_dim"] = state_dict[k].shape[1]
136
+ attention_resolutions.append(attn_res)
137
+ attn_res *= 2
138
+
139
+ match["attention_resolutions"] = attention_resolutions
140
+
141
+ match["model_channels"] = state_dict["conv_in.weight"].shape[0]
142
+ match["in_channels"] = state_dict["conv_in.weight"].shape[1]
143
+ match["adm_in_channels"] = None
144
+ if "class_embedding.linear_1.weight" in state_dict:
145
+ match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1]
146
+ elif "add_embedding.linear_1.weight" in state_dict:
147
+ match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]
148
+
149
+ SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
150
+ 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320,
151
+ 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 2, 10], 'channel_mult': [1, 2, 4],
152
+ 'transformer_depth_middle': 10, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64}
153
+
154
+ SDXL_refiner = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
155
+ 'num_classes': 'sequential', 'adm_in_channels': 2560, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 384,
156
+ 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 4, 4, 0], 'channel_mult': [1, 2, 4, 4],
157
+ 'transformer_depth_middle': 4, 'use_linear_in_transformer': True, 'context_dim': 1280, "num_head_channels": 64}
158
+
159
+ SD21 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
160
+ 'adm_in_channels': None, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2,
161
+ 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4],
162
+ 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024, "num_head_channels": 64}
163
+
164
+ SD21_uncliph = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
165
+ 'num_classes': 'sequential', 'adm_in_channels': 2048, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320,
166
+ 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4],
167
+ 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024, "num_head_channels": 64}
168
+
169
+ SD21_unclipl = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
170
+ 'num_classes': 'sequential', 'adm_in_channels': 1536, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320,
171
+ 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4],
172
+ 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024}
173
+
174
+ SD15 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
175
+ 'adm_in_channels': None, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2,
176
+ 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4],
177
+ 'transformer_depth_middle': 1, 'use_linear_in_transformer': False, 'context_dim': 768, "num_heads": 8}
178
+
179
+ SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
180
+ 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320,
181
+ 'num_res_blocks': 2, 'attention_resolutions': [4], 'transformer_depth': [0, 0, 1], 'channel_mult': [1, 2, 4],
182
+ 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64}
183
+
184
+ SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
185
+ 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320,
186
+ 'num_res_blocks': 2, 'attention_resolutions': [], 'transformer_depth': [0, 0, 0], 'channel_mult': [1, 2, 4],
187
+ 'transformer_depth_middle': 0, 'use_linear_in_transformer': True, "num_head_channels": 64, 'context_dim': 1}
188
+
189
+ SDXL_diffusers_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
190
+ 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 9, 'model_channels': 320,
191
+ 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 2, 10], 'channel_mult': [1, 2, 4],
192
+ 'transformer_depth_middle': 10, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64}
193
+
194
+ supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint]
195
+
196
+ for unet_config in supported_models:
197
+ matches = True
198
+ for k in match:
199
+ if match[k] != unet_config[k]:
200
+ matches = False
201
+ break
202
+ if matches:
203
+ return unet_config
204
+ return None
205
+
206
+ def model_config_from_diffusers_unet(state_dict, use_fp16):
207
+ unet_config = unet_config_from_diffusers_unet(state_dict, use_fp16)
208
+ if unet_config is not None:
209
+ return model_config_from_unet_config(unet_config)
210
+ return None
comfy/model_management.py ADDED
@@ -0,0 +1,711 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import psutil
2
+ from enum import Enum
3
+ from comfy.cli_args import args
4
+ import comfy.utils
5
+ import torch
6
+ import sys
7
+
8
+ class VRAMState(Enum):
9
+ DISABLED = 0 #No vram present: no need to move models to vram
10
+ NO_VRAM = 1 #Very low vram: enable all the options to save vram
11
+ LOW_VRAM = 2
12
+ NORMAL_VRAM = 3
13
+ HIGH_VRAM = 4
14
+ SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.
15
+
16
+ class CPUState(Enum):
17
+ GPU = 0
18
+ CPU = 1
19
+ MPS = 2
20
+
21
+ # Determine VRAM State
22
+ vram_state = VRAMState.NORMAL_VRAM
23
+ set_vram_to = VRAMState.NORMAL_VRAM
24
+ cpu_state = CPUState.GPU
25
+
26
+ total_vram = 0
27
+
28
+ lowvram_available = True
29
+ xpu_available = False
30
+
31
+ directml_enabled = False
32
+ if args.directml is not None:
33
+ import torch_directml
34
+ directml_enabled = True
35
+ device_index = args.directml
36
+ if device_index < 0:
37
+ directml_device = torch_directml.device()
38
+ else:
39
+ directml_device = torch_directml.device(device_index)
40
+ print("Using directml with device:", torch_directml.device_name(device_index))
41
+ # torch_directml.disable_tiled_resources(True)
42
+ lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
43
+
44
+ try:
45
+ import intel_extension_for_pytorch as ipex
46
+ if torch.xpu.is_available():
47
+ xpu_available = True
48
+ except:
49
+ pass
50
+
51
+ try:
52
+ if torch.backends.mps.is_available():
53
+ cpu_state = CPUState.MPS
54
+ import torch.mps
55
+ except:
56
+ pass
57
+
58
+ if args.cpu:
59
+ cpu_state = CPUState.CPU
60
+
61
+ def is_intel_xpu():
62
+ global cpu_state
63
+ global xpu_available
64
+ if cpu_state == CPUState.GPU:
65
+ if xpu_available:
66
+ return True
67
+ return False
68
+
69
+ def get_torch_device():
70
+ global directml_enabled
71
+ global cpu_state
72
+ if directml_enabled:
73
+ global directml_device
74
+ return directml_device
75
+ if cpu_state == CPUState.MPS:
76
+ return torch.device("mps")
77
+ if cpu_state == CPUState.CPU:
78
+ return torch.device("cpu")
79
+ else:
80
+ if is_intel_xpu():
81
+ return torch.device("xpu")
82
+ else:
83
+ return torch.device(torch.cuda.current_device())
84
+
85
+ def get_total_memory(dev=None, torch_total_too=False):
86
+ global directml_enabled
87
+ if dev is None:
88
+ dev = get_torch_device()
89
+
90
+ if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
91
+ mem_total = psutil.virtual_memory().total
92
+ mem_total_torch = mem_total
93
+ else:
94
+ if directml_enabled:
95
+ mem_total = 1024 * 1024 * 1024 #TODO
96
+ mem_total_torch = mem_total
97
+ elif is_intel_xpu():
98
+ stats = torch.xpu.memory_stats(dev)
99
+ mem_reserved = stats['reserved_bytes.all.current']
100
+ mem_total = torch.xpu.get_device_properties(dev).total_memory
101
+ mem_total_torch = mem_reserved
102
+ else:
103
+ stats = torch.cuda.memory_stats(dev)
104
+ mem_reserved = stats['reserved_bytes.all.current']
105
+ _, mem_total_cuda = torch.cuda.mem_get_info(dev)
106
+ mem_total_torch = mem_reserved
107
+ mem_total = mem_total_cuda
108
+
109
+ if torch_total_too:
110
+ return (mem_total, mem_total_torch)
111
+ else:
112
+ return mem_total
113
+
114
+ total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
115
+ total_ram = psutil.virtual_memory().total / (1024 * 1024)
116
+ print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
117
+ if not args.normalvram and not args.cpu:
118
+ if lowvram_available and total_vram <= 4096:
119
+ print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
120
+ set_vram_to = VRAMState.LOW_VRAM
121
+
122
+ try:
123
+ OOM_EXCEPTION = torch.cuda.OutOfMemoryError
124
+ except:
125
+ OOM_EXCEPTION = Exception
126
+
127
+ XFORMERS_VERSION = ""
128
+ XFORMERS_ENABLED_VAE = True
129
+ if args.disable_xformers:
130
+ XFORMERS_IS_AVAILABLE = False
131
+ else:
132
+ try:
133
+ import xformers
134
+ import xformers.ops
135
+ XFORMERS_IS_AVAILABLE = True
136
+ try:
137
+ XFORMERS_VERSION = xformers.version.__version__
138
+ print("xformers version:", XFORMERS_VERSION)
139
+ if XFORMERS_VERSION.startswith("0.0.18"):
140
+ print()
141
+ print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
142
+ print("Please downgrade or upgrade xformers to a different version.")
143
+ print()
144
+ XFORMERS_ENABLED_VAE = False
145
+ except:
146
+ pass
147
+ except:
148
+ XFORMERS_IS_AVAILABLE = False
149
+
150
+ def is_nvidia():
151
+ global cpu_state
152
+ if cpu_state == CPUState.GPU:
153
+ if torch.version.cuda:
154
+ return True
155
+ return False
156
+
157
+ ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
158
+ VAE_DTYPE = torch.float32
159
+
160
+ try:
161
+ if is_nvidia():
162
+ torch_version = torch.version.__version__
163
+ if int(torch_version[0]) >= 2:
164
+ if ENABLE_PYTORCH_ATTENTION == False and XFORMERS_IS_AVAILABLE == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
165
+ ENABLE_PYTORCH_ATTENTION = True
166
+ if torch.cuda.is_bf16_supported():
167
+ VAE_DTYPE = torch.bfloat16
168
+ if is_intel_xpu():
169
+ if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
170
+ ENABLE_PYTORCH_ATTENTION = True
171
+ except:
172
+ pass
173
+
174
+ if is_intel_xpu():
175
+ VAE_DTYPE = torch.bfloat16
176
+
177
+ if args.fp16_vae:
178
+ VAE_DTYPE = torch.float16
179
+ elif args.bf16_vae:
180
+ VAE_DTYPE = torch.bfloat16
181
+ elif args.fp32_vae:
182
+ VAE_DTYPE = torch.float32
183
+
184
+
185
+ if ENABLE_PYTORCH_ATTENTION:
186
+ torch.backends.cuda.enable_math_sdp(True)
187
+ torch.backends.cuda.enable_flash_sdp(True)
188
+ torch.backends.cuda.enable_mem_efficient_sdp(True)
189
+ XFORMERS_IS_AVAILABLE = False
190
+
191
+ if args.lowvram:
192
+ set_vram_to = VRAMState.LOW_VRAM
193
+ lowvram_available = True
194
+ elif args.novram:
195
+ set_vram_to = VRAMState.NO_VRAM
196
+ elif args.highvram or args.gpu_only:
197
+ vram_state = VRAMState.HIGH_VRAM
198
+
199
+ FORCE_FP32 = False
200
+ FORCE_FP16 = False
201
+ if args.force_fp32:
202
+ print("Forcing FP32, if this improves things please report it.")
203
+ FORCE_FP32 = True
204
+
205
+ if args.force_fp16:
206
+ print("Forcing FP16.")
207
+ FORCE_FP16 = True
208
+
209
+ if lowvram_available:
210
+ try:
211
+ import accelerate
212
+ if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
213
+ vram_state = set_vram_to
214
+ except Exception as e:
215
+ import traceback
216
+ print(traceback.format_exc())
217
+ print("ERROR: LOW VRAM MODE NEEDS accelerate.")
218
+ lowvram_available = False
219
+
220
+
221
+ if cpu_state != CPUState.GPU:
222
+ vram_state = VRAMState.DISABLED
223
+
224
+ if cpu_state == CPUState.MPS:
225
+ vram_state = VRAMState.SHARED
226
+
227
+ print(f"Set vram state to: {vram_state.name}")
228
+
229
+ DISABLE_SMART_MEMORY = args.disable_smart_memory
230
+
231
+ if DISABLE_SMART_MEMORY:
232
+ print("Disabling smart memory management")
233
+
234
+ def get_torch_device_name(device):
235
+ if hasattr(device, 'type'):
236
+ if device.type == "cuda":
237
+ try:
238
+ allocator_backend = torch.cuda.get_allocator_backend()
239
+ except:
240
+ allocator_backend = ""
241
+ return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
242
+ else:
243
+ return "{}".format(device.type)
244
+ elif is_intel_xpu():
245
+ return "{} {}".format(device, torch.xpu.get_device_name(device))
246
+ else:
247
+ return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
248
+
249
+ try:
250
+ print("Device:", get_torch_device_name(get_torch_device()))
251
+ except:
252
+ print("Could not pick default device.")
253
+
254
+ print("VAE dtype:", VAE_DTYPE)
255
+
256
+ current_loaded_models = []
257
+
258
+ class LoadedModel:
259
+ def __init__(self, model):
260
+ self.model = model
261
+ self.model_accelerated = False
262
+ self.device = model.load_device
263
+
264
+ def model_memory(self):
265
+ return self.model.model_size()
266
+
267
+ def model_memory_required(self, device):
268
+ if device == self.model.current_device:
269
+ return 0
270
+ else:
271
+ return self.model_memory()
272
+
273
+ def model_load(self, lowvram_model_memory=0):
274
+ patch_model_to = None
275
+ if lowvram_model_memory == 0:
276
+ patch_model_to = self.device
277
+
278
+ self.model.model_patches_to(self.device)
279
+ self.model.model_patches_to(self.model.model_dtype())
280
+
281
+ try:
282
+ self.real_model = self.model.patch_model(device_to=patch_model_to) #TODO: do something with loras and offloading to CPU
283
+ except Exception as e:
284
+ self.model.unpatch_model(self.model.offload_device)
285
+ self.model_unload()
286
+ raise e
287
+
288
+ if lowvram_model_memory > 0:
289
+ print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
290
+ device_map = accelerate.infer_auto_device_map(self.real_model, max_memory={0: "{}MiB".format(lowvram_model_memory // (1024 * 1024)), "cpu": "16GiB"})
291
+ accelerate.dispatch_model(self.real_model, device_map=device_map, main_device=self.device)
292
+ self.model_accelerated = True
293
+
294
+ if is_intel_xpu() and not args.disable_ipex_optimize:
295
+ self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
296
+
297
+ return self.real_model
298
+
299
+ def model_unload(self):
300
+ if self.model_accelerated:
301
+ accelerate.hooks.remove_hook_from_submodules(self.real_model)
302
+ self.model_accelerated = False
303
+
304
+ self.model.unpatch_model(self.model.offload_device)
305
+ self.model.model_patches_to(self.model.offload_device)
306
+
307
+ def __eq__(self, other):
308
+ return self.model is other.model
309
+
310
+ def minimum_inference_memory():
311
+ return (1024 * 1024 * 1024)
312
+
313
+ def unload_model_clones(model):
314
+ to_unload = []
315
+ for i in range(len(current_loaded_models)):
316
+ if model.is_clone(current_loaded_models[i].model):
317
+ to_unload = [i] + to_unload
318
+
319
+ for i in to_unload:
320
+ print("unload clone", i)
321
+ current_loaded_models.pop(i).model_unload()
322
+
323
+ def free_memory(memory_required, device, keep_loaded=[]):
324
+ unloaded_model = False
325
+ for i in range(len(current_loaded_models) -1, -1, -1):
326
+ if not DISABLE_SMART_MEMORY:
327
+ if get_free_memory(device) > memory_required:
328
+ break
329
+ shift_model = current_loaded_models[i]
330
+ if shift_model.device == device:
331
+ if shift_model not in keep_loaded:
332
+ m = current_loaded_models.pop(i)
333
+ m.model_unload()
334
+ del m
335
+ unloaded_model = True
336
+
337
+ if unloaded_model:
338
+ soft_empty_cache()
339
+
340
+
341
+ def load_models_gpu(models, memory_required=0):
342
+ global vram_state
343
+
344
+ inference_memory = minimum_inference_memory()
345
+ extra_mem = max(inference_memory, memory_required)
346
+
347
+ models_to_load = []
348
+ models_already_loaded = []
349
+ for x in models:
350
+ loaded_model = LoadedModel(x)
351
+
352
+ if loaded_model in current_loaded_models:
353
+ index = current_loaded_models.index(loaded_model)
354
+ current_loaded_models.insert(0, current_loaded_models.pop(index))
355
+ models_already_loaded.append(loaded_model)
356
+ else:
357
+ models_to_load.append(loaded_model)
358
+
359
+ if len(models_to_load) == 0:
360
+ devs = set(map(lambda a: a.device, models_already_loaded))
361
+ for d in devs:
362
+ if d != torch.device("cpu"):
363
+ free_memory(extra_mem, d, models_already_loaded)
364
+ return
365
+
366
+ print("loading new")
367
+
368
+ total_memory_required = {}
369
+ for loaded_model in models_to_load:
370
+ unload_model_clones(loaded_model.model)
371
+ total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
372
+
373
+ for device in total_memory_required:
374
+ if device != torch.device("cpu"):
375
+ free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
376
+
377
+ for loaded_model in models_to_load:
378
+ model = loaded_model.model
379
+ torch_dev = model.load_device
380
+ if is_device_cpu(torch_dev):
381
+ vram_set_state = VRAMState.DISABLED
382
+ else:
383
+ vram_set_state = vram_state
384
+ lowvram_model_memory = 0
385
+ if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
386
+ model_size = loaded_model.model_memory_required(torch_dev)
387
+ current_free_mem = get_free_memory(torch_dev)
388
+ lowvram_model_memory = int(max(256 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
389
+ if model_size > (current_free_mem - inference_memory): #only switch to lowvram if really necessary
390
+ vram_set_state = VRAMState.LOW_VRAM
391
+ else:
392
+ lowvram_model_memory = 0
393
+
394
+ if vram_set_state == VRAMState.NO_VRAM:
395
+ lowvram_model_memory = 256 * 1024 * 1024
396
+
397
+ cur_loaded_model = loaded_model.model_load(lowvram_model_memory)
398
+ current_loaded_models.insert(0, loaded_model)
399
+ return
400
+
401
+
402
+ def load_model_gpu(model):
403
+ return load_models_gpu([model])
404
+
405
+ def cleanup_models():
406
+ to_delete = []
407
+ for i in range(len(current_loaded_models)):
408
+ print(sys.getrefcount(current_loaded_models[i].model))
409
+ if sys.getrefcount(current_loaded_models[i].model) <= 2:
410
+ to_delete = [i] + to_delete
411
+
412
+ for i in to_delete:
413
+ x = current_loaded_models.pop(i)
414
+ x.model_unload()
415
+ del x
416
+
417
+ def dtype_size(dtype):
418
+ dtype_size = 4
419
+ if dtype == torch.float16 or dtype == torch.bfloat16:
420
+ dtype_size = 2
421
+ return dtype_size
422
+
423
+ def unet_offload_device():
424
+ if vram_state == VRAMState.HIGH_VRAM:
425
+ return get_torch_device()
426
+ else:
427
+ return torch.device("cpu")
428
+
429
+ def unet_inital_load_device(parameters, dtype):
430
+ torch_dev = get_torch_device()
431
+ if vram_state == VRAMState.HIGH_VRAM:
432
+ return torch_dev
433
+
434
+ cpu_dev = torch.device("cpu")
435
+ if DISABLE_SMART_MEMORY:
436
+ return cpu_dev
437
+
438
+ model_size = dtype_size(dtype) * parameters
439
+
440
+ mem_dev = get_free_memory(torch_dev)
441
+ mem_cpu = get_free_memory(cpu_dev)
442
+ if mem_dev > mem_cpu and model_size < mem_dev:
443
+ return torch_dev
444
+ else:
445
+ return cpu_dev
446
+
447
+ def text_encoder_offload_device():
448
+ if args.gpu_only:
449
+ return get_torch_device()
450
+ else:
451
+ return torch.device("cpu")
452
+
453
+ def text_encoder_device():
454
+ if args.gpu_only:
455
+ return get_torch_device()
456
+ elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
457
+ if is_intel_xpu():
458
+ return torch.device("cpu")
459
+ if should_use_fp16(prioritize_performance=False):
460
+ return get_torch_device()
461
+ else:
462
+ return torch.device("cpu")
463
+ else:
464
+ return torch.device("cpu")
465
+
466
+ def vae_device():
467
+ return get_torch_device()
468
+
469
+ def vae_offload_device():
470
+ if args.gpu_only:
471
+ return get_torch_device()
472
+ else:
473
+ return torch.device("cpu")
474
+
475
+ def vae_dtype():
476
+ global VAE_DTYPE
477
+ return VAE_DTYPE
478
+
479
+ def get_autocast_device(dev):
480
+ if hasattr(dev, 'type'):
481
+ return dev.type
482
+ return "cuda"
483
+
484
+ def cast_to_device(tensor, device, dtype, copy=False):
485
+ device_supports_cast = False
486
+ if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
487
+ device_supports_cast = True
488
+ elif tensor.dtype == torch.bfloat16:
489
+ if hasattr(device, 'type') and device.type.startswith("cuda"):
490
+ device_supports_cast = True
491
+ elif is_intel_xpu():
492
+ device_supports_cast = True
493
+
494
+ if device_supports_cast:
495
+ if copy:
496
+ if tensor.device == device:
497
+ return tensor.to(dtype, copy=copy)
498
+ return tensor.to(device, copy=copy).to(dtype)
499
+ else:
500
+ return tensor.to(device).to(dtype)
501
+ else:
502
+ return tensor.to(dtype).to(device, copy=copy)
503
+
504
+ def xformers_enabled():
505
+ global directml_enabled
506
+ global cpu_state
507
+ if cpu_state != CPUState.GPU:
508
+ return False
509
+ if is_intel_xpu():
510
+ return False
511
+ if directml_enabled:
512
+ return False
513
+ return XFORMERS_IS_AVAILABLE
514
+
515
+
516
+ def xformers_enabled_vae():
517
+ enabled = xformers_enabled()
518
+ if not enabled:
519
+ return False
520
+
521
+ return XFORMERS_ENABLED_VAE
522
+
523
+ def pytorch_attention_enabled():
524
+ global ENABLE_PYTORCH_ATTENTION
525
+ return ENABLE_PYTORCH_ATTENTION
526
+
527
+ def pytorch_attention_flash_attention():
528
+ global ENABLE_PYTORCH_ATTENTION
529
+ if ENABLE_PYTORCH_ATTENTION:
530
+ #TODO: more reliable way of checking for flash attention?
531
+ if is_nvidia(): #pytorch flash attention only works on Nvidia
532
+ return True
533
+ return False
534
+
535
+ def get_free_memory(dev=None, torch_free_too=False):
536
+ global directml_enabled
537
+ if dev is None:
538
+ dev = get_torch_device()
539
+
540
+ if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
541
+ mem_free_total = psutil.virtual_memory().available
542
+ mem_free_torch = mem_free_total
543
+ else:
544
+ if directml_enabled:
545
+ mem_free_total = 1024 * 1024 * 1024 #TODO
546
+ mem_free_torch = mem_free_total
547
+ elif is_intel_xpu():
548
+ stats = torch.xpu.memory_stats(dev)
549
+ mem_active = stats['active_bytes.all.current']
550
+ mem_allocated = stats['allocated_bytes.all.current']
551
+ mem_reserved = stats['reserved_bytes.all.current']
552
+ mem_free_torch = mem_reserved - mem_active
553
+ mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
554
+ else:
555
+ stats = torch.cuda.memory_stats(dev)
556
+ mem_active = stats['active_bytes.all.current']
557
+ mem_reserved = stats['reserved_bytes.all.current']
558
+ mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
559
+ mem_free_torch = mem_reserved - mem_active
560
+ mem_free_total = mem_free_cuda + mem_free_torch
561
+
562
+ if torch_free_too:
563
+ return (mem_free_total, mem_free_torch)
564
+ else:
565
+ return mem_free_total
566
+
567
+ def batch_area_memory(area):
568
+ if xformers_enabled() or pytorch_attention_flash_attention():
569
+ #TODO: these formulas are copied from maximum_batch_area below
570
+ return (area / 20) * (1024 * 1024)
571
+ else:
572
+ return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)
573
+
574
+ def maximum_batch_area():
575
+ global vram_state
576
+ if vram_state == VRAMState.NO_VRAM:
577
+ return 0
578
+
579
+ memory_free = get_free_memory() / (1024 * 1024)
580
+ if xformers_enabled() or pytorch_attention_flash_attention():
581
+ #TODO: this needs to be tweaked
582
+ area = 20 * memory_free
583
+ else:
584
+ #TODO: this formula is because AMD sucks and has memory management issues which might be fixed in the future
585
+ area = ((memory_free - 1024) * 0.9) / (0.6)
586
+ return int(max(area, 0))
587
+
588
+ def cpu_mode():
589
+ global cpu_state
590
+ return cpu_state == CPUState.CPU
591
+
592
+ def mps_mode():
593
+ global cpu_state
594
+ return cpu_state == CPUState.MPS
595
+
596
+ def is_device_cpu(device):
597
+ if hasattr(device, 'type'):
598
+ if (device.type == 'cpu'):
599
+ return True
600
+ return False
601
+
602
+ def is_device_mps(device):
603
+ if hasattr(device, 'type'):
604
+ if (device.type == 'mps'):
605
+ return True
606
+ return False
607
+
608
+ def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
609
+ global directml_enabled
610
+
611
+ if device is not None:
612
+ if is_device_cpu(device):
613
+ return False
614
+
615
+ if FORCE_FP16:
616
+ return True
617
+
618
+ if device is not None: #TODO
619
+ if is_device_mps(device):
620
+ return False
621
+
622
+ if FORCE_FP32:
623
+ return False
624
+
625
+ if directml_enabled:
626
+ return False
627
+
628
+ if cpu_mode() or mps_mode():
629
+ return False #TODO ?
630
+
631
+ if is_intel_xpu():
632
+ return True
633
+
634
+ if torch.cuda.is_bf16_supported():
635
+ return True
636
+
637
+ props = torch.cuda.get_device_properties("cuda")
638
+ if props.major < 6:
639
+ return False
640
+
641
+ fp16_works = False
642
+ #FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled
643
+ #when the model doesn't actually fit on the card
644
+ #TODO: actually test if GP106 and others have the same type of behavior
645
+ nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"]
646
+ for x in nvidia_10_series:
647
+ if x in props.name.lower():
648
+ fp16_works = True
649
+
650
+ if fp16_works:
651
+ free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
652
+ if (not prioritize_performance) or model_params * 4 > free_model_memory:
653
+ return True
654
+
655
+ if props.major < 7:
656
+ return False
657
+
658
+ #FP16 is just broken on these cards
659
+ nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX"]
660
+ for x in nvidia_16_series:
661
+ if x in props.name:
662
+ return False
663
+
664
+ return True
665
+
666
+ def soft_empty_cache(force=False):
667
+ global cpu_state
668
+ if cpu_state == CPUState.MPS:
669
+ torch.mps.empty_cache()
670
+ elif is_intel_xpu():
671
+ torch.xpu.empty_cache()
672
+ elif torch.cuda.is_available():
673
+ if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
674
+ torch.cuda.empty_cache()
675
+ torch.cuda.ipc_collect()
676
+
677
+ def resolve_lowvram_weight(weight, model, key):
678
+ if weight.device == torch.device("meta"): #lowvram NOTE: this depends on the inner working of the accelerate library so it might break.
679
+ key_split = key.split('.') # I have no idea why they don't just leave the weight there instead of using the meta device.
680
+ op = comfy.utils.get_attr(model, '.'.join(key_split[:-1]))
681
+ weight = op._hf_hook.weights_map[key_split[-1]]
682
+ return weight
683
+
684
+ #TODO: might be cleaner to put this somewhere else
685
+ import threading
686
+
687
+ class InterruptProcessingException(Exception):
688
+ pass
689
+
690
+ interrupt_processing_mutex = threading.RLock()
691
+
692
+ interrupt_processing = False
693
+ def interrupt_current_processing(value=True):
694
+ global interrupt_processing
695
+ global interrupt_processing_mutex
696
+ with interrupt_processing_mutex:
697
+ interrupt_processing = value
698
+
699
+ def processing_interrupted():
700
+ global interrupt_processing
701
+ global interrupt_processing_mutex
702
+ with interrupt_processing_mutex:
703
+ return interrupt_processing
704
+
705
+ def throw_exception_if_processing_interrupted():
706
+ global interrupt_processing
707
+ global interrupt_processing_mutex
708
+ with interrupt_processing_mutex:
709
+ if interrupt_processing:
710
+ interrupt_processing = False
711
+ raise InterruptProcessingException()
comfy/model_patcher.py ADDED
@@ -0,0 +1,288 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import copy
3
+ import inspect
4
+
5
+ import comfy.utils
6
+ import comfy.model_management
7
+
8
+ class ModelPatcher:
9
+ def __init__(self, model, load_device, offload_device, size=0, current_device=None):
10
+ self.size = size
11
+ self.model = model
12
+ self.patches = {}
13
+ self.backup = {}
14
+ self.model_options = {"transformer_options":{}}
15
+ self.model_size()
16
+ self.load_device = load_device
17
+ self.offload_device = offload_device
18
+ if current_device is None:
19
+ self.current_device = self.offload_device
20
+ else:
21
+ self.current_device = current_device
22
+
23
+ def model_size(self):
24
+ if self.size > 0:
25
+ return self.size
26
+ model_sd = self.model.state_dict()
27
+ size = 0
28
+ for k in model_sd:
29
+ t = model_sd[k]
30
+ size += t.nelement() * t.element_size()
31
+ self.size = size
32
+ self.model_keys = set(model_sd.keys())
33
+ return size
34
+
35
+ def clone(self):
36
+ n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device)
37
+ n.patches = {}
38
+ for k in self.patches:
39
+ n.patches[k] = self.patches[k][:]
40
+
41
+ n.model_options = copy.deepcopy(self.model_options)
42
+ n.model_keys = self.model_keys
43
+ return n
44
+
45
+ def is_clone(self, other):
46
+ if hasattr(other, 'model') and self.model is other.model:
47
+ return True
48
+ return False
49
+
50
+ def set_model_sampler_cfg_function(self, sampler_cfg_function):
51
+ if len(inspect.signature(sampler_cfg_function).parameters) == 3:
52
+ self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
53
+ else:
54
+ self.model_options["sampler_cfg_function"] = sampler_cfg_function
55
+
56
+ def set_model_unet_function_wrapper(self, unet_wrapper_function):
57
+ self.model_options["model_function_wrapper"] = unet_wrapper_function
58
+
59
+ def set_model_patch(self, patch, name):
60
+ to = self.model_options["transformer_options"]
61
+ if "patches" not in to:
62
+ to["patches"] = {}
63
+ to["patches"][name] = to["patches"].get(name, []) + [patch]
64
+
65
+ def set_model_patch_replace(self, patch, name, block_name, number):
66
+ to = self.model_options["transformer_options"]
67
+ if "patches_replace" not in to:
68
+ to["patches_replace"] = {}
69
+ if name not in to["patches_replace"]:
70
+ to["patches_replace"][name] = {}
71
+ to["patches_replace"][name][(block_name, number)] = patch
72
+
73
+ def set_model_attn1_patch(self, patch):
74
+ self.set_model_patch(patch, "attn1_patch")
75
+
76
+ def set_model_attn2_patch(self, patch):
77
+ self.set_model_patch(patch, "attn2_patch")
78
+
79
+ def set_model_attn1_replace(self, patch, block_name, number):
80
+ self.set_model_patch_replace(patch, "attn1", block_name, number)
81
+
82
+ def set_model_attn2_replace(self, patch, block_name, number):
83
+ self.set_model_patch_replace(patch, "attn2", block_name, number)
84
+
85
+ def set_model_attn1_output_patch(self, patch):
86
+ self.set_model_patch(patch, "attn1_output_patch")
87
+
88
+ def set_model_attn2_output_patch(self, patch):
89
+ self.set_model_patch(patch, "attn2_output_patch")
90
+
91
+ def set_model_output_block_patch(self, patch):
92
+ self.set_model_patch(patch, "output_block_patch")
93
+
94
+ def model_patches_to(self, device):
95
+ to = self.model_options["transformer_options"]
96
+ if "patches" in to:
97
+ patches = to["patches"]
98
+ for name in patches:
99
+ patch_list = patches[name]
100
+ for i in range(len(patch_list)):
101
+ if hasattr(patch_list[i], "to"):
102
+ patch_list[i] = patch_list[i].to(device)
103
+ if "patches_replace" in to:
104
+ patches = to["patches_replace"]
105
+ for name in patches:
106
+ patch_list = patches[name]
107
+ for k in patch_list:
108
+ if hasattr(patch_list[k], "to"):
109
+ patch_list[k] = patch_list[k].to(device)
110
+
111
+ def model_dtype(self):
112
+ if hasattr(self.model, "get_dtype"):
113
+ return self.model.get_dtype()
114
+
115
+ def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
116
+ p = set()
117
+ for k in patches:
118
+ if k in self.model_keys:
119
+ p.add(k)
120
+ current_patches = self.patches.get(k, [])
121
+ current_patches.append((strength_patch, patches[k], strength_model))
122
+ self.patches[k] = current_patches
123
+
124
+ return list(p)
125
+
126
+ def get_key_patches(self, filter_prefix=None):
127
+ model_sd = self.model_state_dict()
128
+ p = {}
129
+ for k in model_sd:
130
+ if filter_prefix is not None:
131
+ if not k.startswith(filter_prefix):
132
+ continue
133
+ if k in self.patches:
134
+ p[k] = [model_sd[k]] + self.patches[k]
135
+ else:
136
+ p[k] = (model_sd[k],)
137
+ return p
138
+
139
+ def model_state_dict(self, filter_prefix=None):
140
+ sd = self.model.state_dict()
141
+ keys = list(sd.keys())
142
+ if filter_prefix is not None:
143
+ for k in keys:
144
+ if not k.startswith(filter_prefix):
145
+ sd.pop(k)
146
+ return sd
147
+
148
+ def patch_model(self, device_to=None):
149
+ model_sd = self.model_state_dict()
150
+ for key in self.patches:
151
+ if key not in model_sd:
152
+ print("could not patch. key doesn't exist in model:", key)
153
+ continue
154
+
155
+ weight = model_sd[key]
156
+
157
+ if key not in self.backup:
158
+ self.backup[key] = weight.to(self.offload_device)
159
+
160
+ if device_to is not None:
161
+ temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
162
+ else:
163
+ temp_weight = weight.to(torch.float32, copy=True)
164
+ out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
165
+ comfy.utils.set_attr(self.model, key, out_weight)
166
+ del temp_weight
167
+
168
+ if device_to is not None:
169
+ self.model.to(device_to)
170
+ self.current_device = device_to
171
+
172
+ return self.model
173
+
174
+ def calculate_weight(self, patches, weight, key):
175
+ for p in patches:
176
+ alpha = p[0]
177
+ v = p[1]
178
+ strength_model = p[2]
179
+
180
+ if strength_model != 1.0:
181
+ weight *= strength_model
182
+
183
+ if isinstance(v, list):
184
+ v = (self.calculate_weight(v[1:], v[0].clone(), key), )
185
+
186
+ if len(v) == 1:
187
+ w1 = v[0]
188
+ if alpha != 0.0:
189
+ if w1.shape != weight.shape:
190
+ print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
191
+ else:
192
+ weight += alpha * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype)
193
+ elif len(v) == 4: #lora/locon
194
+ mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32)
195
+ mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32)
196
+ if v[2] is not None:
197
+ alpha *= v[2] / mat2.shape[0]
198
+ if v[3] is not None:
199
+ #locon mid weights, hopefully the math is fine because I didn't properly test it
200
+ mat3 = comfy.model_management.cast_to_device(v[3], weight.device, torch.float32)
201
+ final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
202
+ mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
203
+ try:
204
+ weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype)
205
+ except Exception as e:
206
+ print("ERROR", key, e)
207
+ elif len(v) == 8: #lokr
208
+ w1 = v[0]
209
+ w2 = v[1]
210
+ w1_a = v[3]
211
+ w1_b = v[4]
212
+ w2_a = v[5]
213
+ w2_b = v[6]
214
+ t2 = v[7]
215
+ dim = None
216
+
217
+ if w1 is None:
218
+ dim = w1_b.shape[0]
219
+ w1 = torch.mm(comfy.model_management.cast_to_device(w1_a, weight.device, torch.float32),
220
+ comfy.model_management.cast_to_device(w1_b, weight.device, torch.float32))
221
+ else:
222
+ w1 = comfy.model_management.cast_to_device(w1, weight.device, torch.float32)
223
+
224
+ if w2 is None:
225
+ dim = w2_b.shape[0]
226
+ if t2 is None:
227
+ w2 = torch.mm(comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32),
228
+ comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32))
229
+ else:
230
+ w2 = torch.einsum('i j k l, j r, i p -> p r k l',
231
+ comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
232
+ comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32),
233
+ comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32))
234
+ else:
235
+ w2 = comfy.model_management.cast_to_device(w2, weight.device, torch.float32)
236
+
237
+ if len(w2.shape) == 4:
238
+ w1 = w1.unsqueeze(2).unsqueeze(2)
239
+ if v[2] is not None and dim is not None:
240
+ alpha *= v[2] / dim
241
+
242
+ try:
243
+ weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
244
+ except Exception as e:
245
+ print("ERROR", key, e)
246
+ else: #loha
247
+ w1a = v[0]
248
+ w1b = v[1]
249
+ if v[2] is not None:
250
+ alpha *= v[2] / w1b.shape[0]
251
+ w2a = v[3]
252
+ w2b = v[4]
253
+ if v[5] is not None: #cp decomposition
254
+ t1 = v[5]
255
+ t2 = v[6]
256
+ m1 = torch.einsum('i j k l, j r, i p -> p r k l',
257
+ comfy.model_management.cast_to_device(t1, weight.device, torch.float32),
258
+ comfy.model_management.cast_to_device(w1b, weight.device, torch.float32),
259
+ comfy.model_management.cast_to_device(w1a, weight.device, torch.float32))
260
+
261
+ m2 = torch.einsum('i j k l, j r, i p -> p r k l',
262
+ comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
263
+ comfy.model_management.cast_to_device(w2b, weight.device, torch.float32),
264
+ comfy.model_management.cast_to_device(w2a, weight.device, torch.float32))
265
+ else:
266
+ m1 = torch.mm(comfy.model_management.cast_to_device(w1a, weight.device, torch.float32),
267
+ comfy.model_management.cast_to_device(w1b, weight.device, torch.float32))
268
+ m2 = torch.mm(comfy.model_management.cast_to_device(w2a, weight.device, torch.float32),
269
+ comfy.model_management.cast_to_device(w2b, weight.device, torch.float32))
270
+
271
+ try:
272
+ weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
273
+ except Exception as e:
274
+ print("ERROR", key, e)
275
+
276
+ return weight
277
+
278
+ def unpatch_model(self, device_to=None):
279
+ keys = list(self.backup.keys())
280
+
281
+ for k in keys:
282
+ comfy.utils.set_attr(self.model, k, self.backup[k])
283
+
284
+ self.backup = {}
285
+
286
+ if device_to is not None:
287
+ self.model.to(device_to)
288
+ self.current_device = device_to
comfy/ops.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from contextlib import contextmanager
3
+
4
+ class Linear(torch.nn.Module):
5
+ def __init__(self, in_features: int, out_features: int, bias: bool = True,
6
+ device=None, dtype=None) -> None:
7
+ factory_kwargs = {'device': device, 'dtype': dtype}
8
+ super().__init__()
9
+ self.in_features = in_features
10
+ self.out_features = out_features
11
+ self.weight = torch.nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
12
+ if bias:
13
+ self.bias = torch.nn.Parameter(torch.empty(out_features, **factory_kwargs))
14
+ else:
15
+ self.register_parameter('bias', None)
16
+
17
+ def forward(self, input):
18
+ return torch.nn.functional.linear(input, self.weight, self.bias)
19
+
20
+ class Conv2d(torch.nn.Conv2d):
21
+ def reset_parameters(self):
22
+ return None
23
+
24
+ def conv_nd(dims, *args, **kwargs):
25
+ if dims == 2:
26
+ return Conv2d(*args, **kwargs)
27
+ else:
28
+ raise ValueError(f"unsupported dimensions: {dims}")
29
+
30
+ @contextmanager
31
+ def use_comfy_ops(device=None, dtype=None): # Kind of an ugly hack but I can't think of a better way
32
+ old_torch_nn_linear = torch.nn.Linear
33
+ force_device = device
34
+ force_dtype = dtype
35
+ def linear_with_dtype(in_features: int, out_features: int, bias: bool = True, device=None, dtype=None):
36
+ if force_device is not None:
37
+ device = force_device
38
+ if force_dtype is not None:
39
+ dtype = force_dtype
40
+ return Linear(in_features, out_features, bias=bias, device=device, dtype=dtype)
41
+
42
+ torch.nn.Linear = linear_with_dtype
43
+ try:
44
+ yield
45
+ finally:
46
+ torch.nn.Linear = old_torch_nn_linear
comfy/options.py ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+
2
+ args_parsing = False
3
+
4
+ def enable_args_parsing(enable=True):
5
+ global args_parsing
6
+ args_parsing = enable
comfy/sample.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import comfy.model_management
3
+ import comfy.samplers
4
+ import comfy.utils
5
+ import math
6
+ import numpy as np
7
+
8
+ def prepare_noise(latent_image, seed, noise_inds=None):
9
+ """
10
+ creates random noise given a latent image and a seed.
11
+ optional arg skip can be used to skip and discard x number of noise generations for a given seed
12
+ """
13
+ generator = torch.manual_seed(seed)
14
+ if noise_inds is None:
15
+ return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
16
+
17
+ unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
18
+ noises = []
19
+ for i in range(unique_inds[-1]+1):
20
+ noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
21
+ if i in unique_inds:
22
+ noises.append(noise)
23
+ noises = [noises[i] for i in inverse]
24
+ noises = torch.cat(noises, axis=0)
25
+ return noises
26
+
27
+ def prepare_mask(noise_mask, shape, device):
28
+ """ensures noise mask is of proper dimensions"""
29
+ noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
30
+ noise_mask = noise_mask.round()
31
+ noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
32
+ noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0])
33
+ noise_mask = noise_mask.to(device)
34
+ return noise_mask
35
+
36
+ def broadcast_cond(cond, batch, device):
37
+ """broadcasts conditioning to the batch size"""
38
+ copy = []
39
+ for p in cond:
40
+ t = comfy.utils.repeat_to_batch_size(p[0], batch)
41
+ t = t.to(device)
42
+ copy += [[t] + p[1:]]
43
+ return copy
44
+
45
+ def get_models_from_cond(cond, model_type):
46
+ models = []
47
+ for c in cond:
48
+ if model_type in c[1]:
49
+ models += [c[1][model_type]]
50
+ return models
51
+
52
+ def get_additional_models(positive, negative, dtype):
53
+ """loads additional models in positive and negative conditioning"""
54
+ control_nets = set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control"))
55
+
56
+ inference_memory = 0
57
+ control_models = []
58
+ for m in control_nets:
59
+ control_models += m.get_models()
60
+ inference_memory += m.inference_memory_requirements(dtype)
61
+
62
+ gligen = get_models_from_cond(positive, "gligen") + get_models_from_cond(negative, "gligen")
63
+ gligen = [x[1] for x in gligen]
64
+ models = control_models + gligen
65
+ return models, inference_memory
66
+
67
+ def cleanup_additional_models(models):
68
+ """cleanup additional models that were loaded"""
69
+ for m in models:
70
+ if hasattr(m, 'cleanup'):
71
+ m.cleanup()
72
+
73
+ def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
74
+ device = comfy.model_management.get_torch_device()
75
+
76
+ if noise_mask is not None:
77
+ noise_mask = prepare_mask(noise_mask, noise.shape, device)
78
+
79
+ real_model = None
80
+ models, inference_memory = get_additional_models(positive, negative, model.model_dtype())
81
+ comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]) + inference_memory)
82
+ real_model = model.model
83
+
84
+ noise = noise.to(device)
85
+ latent_image = latent_image.to(device)
86
+
87
+ positive_copy = broadcast_cond(positive, noise.shape[0], device)
88
+ negative_copy = broadcast_cond(negative, noise.shape[0], device)
89
+
90
+
91
+ sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
92
+
93
+ samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
94
+ samples = samples.cpu()
95
+
96
+ cleanup_additional_models(models)
97
+ return samples
comfy/samplers.py ADDED
@@ -0,0 +1,744 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from .k_diffusion import sampling as k_diffusion_sampling
2
+ from .k_diffusion import external as k_diffusion_external
3
+ from .extra_samplers import uni_pc
4
+ import torch
5
+ from comfy import model_management
6
+ from .ldm.models.diffusion.ddim import DDIMSampler
7
+ from .ldm.modules.diffusionmodules.util import make_ddim_timesteps
8
+ import math
9
+ from comfy import model_base
10
+ import comfy.utils
11
+
12
+ def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9)
13
+ return abs(a*b) // math.gcd(a, b)
14
+
15
+ #The main sampling function shared by all the samplers
16
+ #Returns predicted noise
17
+ def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, cond_concat=None, model_options={}, seed=None):
18
+ def get_area_and_mult(cond, x_in, cond_concat_in, timestep_in):
19
+ area = (x_in.shape[2], x_in.shape[3], 0, 0)
20
+ strength = 1.0
21
+ if 'timestep_start' in cond[1]:
22
+ timestep_start = cond[1]['timestep_start']
23
+ if timestep_in[0] > timestep_start:
24
+ return None
25
+ if 'timestep_end' in cond[1]:
26
+ timestep_end = cond[1]['timestep_end']
27
+ if timestep_in[0] < timestep_end:
28
+ return None
29
+ if 'area' in cond[1]:
30
+ area = cond[1]['area']
31
+ if 'strength' in cond[1]:
32
+ strength = cond[1]['strength']
33
+
34
+ adm_cond = None
35
+ if 'adm_encoded' in cond[1]:
36
+ adm_cond = cond[1]['adm_encoded']
37
+
38
+ input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
39
+ if 'mask' in cond[1]:
40
+ # Scale the mask to the size of the input
41
+ # The mask should have been resized as we began the sampling process
42
+ mask_strength = 1.0
43
+ if "mask_strength" in cond[1]:
44
+ mask_strength = cond[1]["mask_strength"]
45
+ mask = cond[1]['mask']
46
+ assert(mask.shape[1] == x_in.shape[2])
47
+ assert(mask.shape[2] == x_in.shape[3])
48
+ mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
49
+ mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
50
+ else:
51
+ mask = torch.ones_like(input_x)
52
+ mult = mask * strength
53
+
54
+ if 'mask' not in cond[1]:
55
+ rr = 8
56
+ if area[2] != 0:
57
+ for t in range(rr):
58
+ mult[:,:,t:1+t,:] *= ((1.0/rr) * (t + 1))
59
+ if (area[0] + area[2]) < x_in.shape[2]:
60
+ for t in range(rr):
61
+ mult[:,:,area[0] - 1 - t:area[0] - t,:] *= ((1.0/rr) * (t + 1))
62
+ if area[3] != 0:
63
+ for t in range(rr):
64
+ mult[:,:,:,t:1+t] *= ((1.0/rr) * (t + 1))
65
+ if (area[1] + area[3]) < x_in.shape[3]:
66
+ for t in range(rr):
67
+ mult[:,:,:,area[1] - 1 - t:area[1] - t] *= ((1.0/rr) * (t + 1))
68
+
69
+ conditionning = {}
70
+ conditionning['c_crossattn'] = cond[0]
71
+ if cond_concat_in is not None and len(cond_concat_in) > 0:
72
+ cropped = []
73
+ for x in cond_concat_in:
74
+ cr = x[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
75
+ cropped.append(cr)
76
+ conditionning['c_concat'] = torch.cat(cropped, dim=1)
77
+
78
+ if adm_cond is not None:
79
+ conditionning['c_adm'] = adm_cond
80
+
81
+ control = None
82
+ if 'control' in cond[1]:
83
+ control = cond[1]['control']
84
+
85
+ patches = None
86
+ if 'gligen' in cond[1]:
87
+ gligen = cond[1]['gligen']
88
+ patches = {}
89
+ gligen_type = gligen[0]
90
+ gligen_model = gligen[1]
91
+ if gligen_type == "position":
92
+ gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
93
+ else:
94
+ gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)
95
+
96
+ patches['middle_patch'] = [gligen_patch]
97
+
98
+ return (input_x, mult, conditionning, area, control, patches)
99
+
100
+ def cond_equal_size(c1, c2):
101
+ if c1 is c2:
102
+ return True
103
+ if c1.keys() != c2.keys():
104
+ return False
105
+ if 'c_crossattn' in c1:
106
+ s1 = c1['c_crossattn'].shape
107
+ s2 = c2['c_crossattn'].shape
108
+ if s1 != s2:
109
+ if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen
110
+ return False
111
+
112
+ mult_min = lcm(s1[1], s2[1])
113
+ diff = mult_min // min(s1[1], s2[1])
114
+ if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
115
+ return False
116
+ if 'c_concat' in c1:
117
+ if c1['c_concat'].shape != c2['c_concat'].shape:
118
+ return False
119
+ if 'c_adm' in c1:
120
+ if c1['c_adm'].shape != c2['c_adm'].shape:
121
+ return False
122
+ return True
123
+
124
+ def can_concat_cond(c1, c2):
125
+ if c1[0].shape != c2[0].shape:
126
+ return False
127
+
128
+ #control
129
+ if (c1[4] is None) != (c2[4] is None):
130
+ return False
131
+ if c1[4] is not None:
132
+ if c1[4] is not c2[4]:
133
+ return False
134
+
135
+ #patches
136
+ if (c1[5] is None) != (c2[5] is None):
137
+ return False
138
+ if (c1[5] is not None):
139
+ if c1[5] is not c2[5]:
140
+ return False
141
+
142
+ return cond_equal_size(c1[2], c2[2])
143
+
144
+ def cond_cat(c_list):
145
+ c_crossattn = []
146
+ c_concat = []
147
+ c_adm = []
148
+ crossattn_max_len = 0
149
+ for x in c_list:
150
+ if 'c_crossattn' in x:
151
+ c = x['c_crossattn']
152
+ if crossattn_max_len == 0:
153
+ crossattn_max_len = c.shape[1]
154
+ else:
155
+ crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
156
+ c_crossattn.append(c)
157
+ if 'c_concat' in x:
158
+ c_concat.append(x['c_concat'])
159
+ if 'c_adm' in x:
160
+ c_adm.append(x['c_adm'])
161
+ out = {}
162
+ c_crossattn_out = []
163
+ for c in c_crossattn:
164
+ if c.shape[1] < crossattn_max_len:
165
+ c = c.repeat(1, crossattn_max_len // c.shape[1], 1) #padding with repeat doesn't change result
166
+ c_crossattn_out.append(c)
167
+
168
+ if len(c_crossattn_out) > 0:
169
+ out['c_crossattn'] = torch.cat(c_crossattn_out)
170
+ if len(c_concat) > 0:
171
+ out['c_concat'] = torch.cat(c_concat)
172
+ if len(c_adm) > 0:
173
+ out['c_adm'] = torch.cat(c_adm)
174
+ return out
175
+
176
+ def calc_cond_uncond_batch(model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in, model_options):
177
+ out_cond = torch.zeros_like(x_in)
178
+ out_count = torch.ones_like(x_in)/100000.0
179
+
180
+ out_uncond = torch.zeros_like(x_in)
181
+ out_uncond_count = torch.ones_like(x_in)/100000.0
182
+
183
+ COND = 0
184
+ UNCOND = 1
185
+
186
+ to_run = []
187
+ for x in cond:
188
+ p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
189
+ if p is None:
190
+ continue
191
+
192
+ to_run += [(p, COND)]
193
+ if uncond is not None:
194
+ for x in uncond:
195
+ p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
196
+ if p is None:
197
+ continue
198
+
199
+ to_run += [(p, UNCOND)]
200
+
201
+ while len(to_run) > 0:
202
+ first = to_run[0]
203
+ first_shape = first[0][0].shape
204
+ to_batch_temp = []
205
+ for x in range(len(to_run)):
206
+ if can_concat_cond(to_run[x][0], first[0]):
207
+ to_batch_temp += [x]
208
+
209
+ to_batch_temp.reverse()
210
+ to_batch = to_batch_temp[:1]
211
+
212
+ for i in range(1, len(to_batch_temp) + 1):
213
+ batch_amount = to_batch_temp[:len(to_batch_temp)//i]
214
+ if (len(batch_amount) * first_shape[0] * first_shape[2] * first_shape[3] < max_total_area):
215
+ to_batch = batch_amount
216
+ break
217
+
218
+ input_x = []
219
+ mult = []
220
+ c = []
221
+ cond_or_uncond = []
222
+ area = []
223
+ control = None
224
+ patches = None
225
+ for x in to_batch:
226
+ o = to_run.pop(x)
227
+ p = o[0]
228
+ input_x += [p[0]]
229
+ mult += [p[1]]
230
+ c += [p[2]]
231
+ area += [p[3]]
232
+ cond_or_uncond += [o[1]]
233
+ control = p[4]
234
+ patches = p[5]
235
+
236
+ batch_chunks = len(cond_or_uncond)
237
+ input_x = torch.cat(input_x)
238
+ c = cond_cat(c)
239
+ timestep_ = torch.cat([timestep] * batch_chunks)
240
+
241
+ if control is not None:
242
+ c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
243
+
244
+ transformer_options = {}
245
+ if 'transformer_options' in model_options:
246
+ transformer_options = model_options['transformer_options'].copy()
247
+
248
+ if patches is not None:
249
+ if "patches" in transformer_options:
250
+ cur_patches = transformer_options["patches"].copy()
251
+ for p in patches:
252
+ if p in cur_patches:
253
+ cur_patches[p] = cur_patches[p] + patches[p]
254
+ else:
255
+ cur_patches[p] = patches[p]
256
+ else:
257
+ transformer_options["patches"] = patches
258
+
259
+ transformer_options["cond_or_uncond"] = cond_or_uncond[:]
260
+ c['transformer_options'] = transformer_options
261
+
262
+ if 'model_function_wrapper' in model_options:
263
+ output = model_options['model_function_wrapper'](model_function, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
264
+ else:
265
+ output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
266
+ del input_x
267
+
268
+ for o in range(batch_chunks):
269
+ if cond_or_uncond[o] == COND:
270
+ out_cond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
271
+ out_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
272
+ else:
273
+ out_uncond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
274
+ out_uncond_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
275
+ del mult
276
+
277
+ out_cond /= out_count
278
+ del out_count
279
+ out_uncond /= out_uncond_count
280
+ del out_uncond_count
281
+
282
+ return out_cond, out_uncond
283
+
284
+
285
+ max_total_area = model_management.maximum_batch_area()
286
+ if math.isclose(cond_scale, 1.0):
287
+ uncond = None
288
+
289
+ cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options)
290
+ if "sampler_cfg_function" in model_options:
291
+ args = {"cond": cond, "uncond": uncond, "cond_scale": cond_scale, "timestep": timestep}
292
+ return model_options["sampler_cfg_function"](args)
293
+ else:
294
+ return uncond + (cond - uncond) * cond_scale
295
+
296
+
297
+ class CompVisVDenoiser(k_diffusion_external.DiscreteVDDPMDenoiser):
298
+ def __init__(self, model, quantize=False, device='cpu'):
299
+ super().__init__(model, model.alphas_cumprod, quantize=quantize)
300
+
301
+ def get_v(self, x, t, cond, **kwargs):
302
+ return self.inner_model.apply_model(x, t, cond, **kwargs)
303
+
304
+
305
+ class CFGNoisePredictor(torch.nn.Module):
306
+ def __init__(self, model):
307
+ super().__init__()
308
+ self.inner_model = model
309
+ self.alphas_cumprod = model.alphas_cumprod
310
+ def apply_model(self, x, timestep, cond, uncond, cond_scale, cond_concat=None, model_options={}, seed=None):
311
+ out = sampling_function(self.inner_model.apply_model, x, timestep, uncond, cond, cond_scale, cond_concat, model_options=model_options, seed=seed)
312
+ return out
313
+
314
+
315
+ class KSamplerX0Inpaint(torch.nn.Module):
316
+ def __init__(self, model):
317
+ super().__init__()
318
+ self.inner_model = model
319
+ def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, cond_concat=None, model_options={}, seed=None):
320
+ if denoise_mask is not None:
321
+ latent_mask = 1. - denoise_mask
322
+ x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
323
+ out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, cond_concat=cond_concat, model_options=model_options, seed=seed)
324
+ if denoise_mask is not None:
325
+ out *= denoise_mask
326
+
327
+ if denoise_mask is not None:
328
+ out += self.latent_image * latent_mask
329
+ return out
330
+
331
+ def simple_scheduler(model, steps):
332
+ sigs = []
333
+ ss = len(model.sigmas) / steps
334
+ for x in range(steps):
335
+ sigs += [float(model.sigmas[-(1 + int(x * ss))])]
336
+ sigs += [0.0]
337
+ return torch.FloatTensor(sigs)
338
+
339
+ def ddim_scheduler(model, steps):
340
+ sigs = []
341
+ ddim_timesteps = make_ddim_timesteps(ddim_discr_method="uniform", num_ddim_timesteps=steps, num_ddpm_timesteps=model.inner_model.inner_model.num_timesteps, verbose=False)
342
+ for x in range(len(ddim_timesteps) - 1, -1, -1):
343
+ ts = ddim_timesteps[x]
344
+ if ts > 999:
345
+ ts = 999
346
+ sigs.append(model.t_to_sigma(torch.tensor(ts)))
347
+ sigs += [0.0]
348
+ return torch.FloatTensor(sigs)
349
+
350
+ def sgm_scheduler(model, steps):
351
+ sigs = []
352
+ timesteps = torch.linspace(model.inner_model.inner_model.num_timesteps - 1, 0, steps + 1)[:-1].type(torch.int)
353
+ for x in range(len(timesteps)):
354
+ ts = timesteps[x]
355
+ if ts > 999:
356
+ ts = 999
357
+ sigs.append(model.t_to_sigma(torch.tensor(ts)))
358
+ sigs += [0.0]
359
+ return torch.FloatTensor(sigs)
360
+
361
+ def blank_inpaint_image_like(latent_image):
362
+ blank_image = torch.ones_like(latent_image)
363
+ # these are the values for "zero" in pixel space translated to latent space
364
+ blank_image[:,0] *= 0.8223
365
+ blank_image[:,1] *= -0.6876
366
+ blank_image[:,2] *= 0.6364
367
+ blank_image[:,3] *= 0.1380
368
+ return blank_image
369
+
370
+ def get_mask_aabb(masks):
371
+ if masks.numel() == 0:
372
+ return torch.zeros((0, 4), device=masks.device, dtype=torch.int)
373
+
374
+ b = masks.shape[0]
375
+
376
+ bounding_boxes = torch.zeros((b, 4), device=masks.device, dtype=torch.int)
377
+ is_empty = torch.zeros((b), device=masks.device, dtype=torch.bool)
378
+ for i in range(b):
379
+ mask = masks[i]
380
+ if mask.numel() == 0:
381
+ continue
382
+ if torch.max(mask != 0) == False:
383
+ is_empty[i] = True
384
+ continue
385
+ y, x = torch.where(mask)
386
+ bounding_boxes[i, 0] = torch.min(x)
387
+ bounding_boxes[i, 1] = torch.min(y)
388
+ bounding_boxes[i, 2] = torch.max(x)
389
+ bounding_boxes[i, 3] = torch.max(y)
390
+
391
+ return bounding_boxes, is_empty
392
+
393
+ def resolve_areas_and_cond_masks(conditions, h, w, device):
394
+ # We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes.
395
+ # While we're doing this, we can also resolve the mask device and scaling for performance reasons
396
+ for i in range(len(conditions)):
397
+ c = conditions[i]
398
+ if 'area' in c[1]:
399
+ area = c[1]['area']
400
+ if area[0] == "percentage":
401
+ modified = c[1].copy()
402
+ area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
403
+ modified['area'] = area
404
+ c = [c[0], modified]
405
+ conditions[i] = c
406
+
407
+ if 'mask' in c[1]:
408
+ mask = c[1]['mask']
409
+ mask = mask.to(device=device)
410
+ modified = c[1].copy()
411
+ if len(mask.shape) == 2:
412
+ mask = mask.unsqueeze(0)
413
+ if mask.shape[1] != h or mask.shape[2] != w:
414
+ mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1)
415
+
416
+ if modified.get("set_area_to_bounds", False):
417
+ bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
418
+ boxes, is_empty = get_mask_aabb(bounds)
419
+ if is_empty[0]:
420
+ # Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway)
421
+ modified['area'] = (8, 8, 0, 0)
422
+ else:
423
+ box = boxes[0]
424
+ H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
425
+ H = max(8, H)
426
+ W = max(8, W)
427
+ area = (int(H), int(W), int(Y), int(X))
428
+ modified['area'] = area
429
+
430
+ modified['mask'] = mask
431
+ conditions[i] = [c[0], modified]
432
+
433
+ def create_cond_with_same_area_if_none(conds, c):
434
+ if 'area' not in c[1]:
435
+ return
436
+
437
+ c_area = c[1]['area']
438
+ smallest = None
439
+ for x in conds:
440
+ if 'area' in x[1]:
441
+ a = x[1]['area']
442
+ if c_area[2] >= a[2] and c_area[3] >= a[3]:
443
+ if a[0] + a[2] >= c_area[0] + c_area[2]:
444
+ if a[1] + a[3] >= c_area[1] + c_area[3]:
445
+ if smallest is None:
446
+ smallest = x
447
+ elif 'area' not in smallest[1]:
448
+ smallest = x
449
+ else:
450
+ if smallest[1]['area'][0] * smallest[1]['area'][1] > a[0] * a[1]:
451
+ smallest = x
452
+ else:
453
+ if smallest is None:
454
+ smallest = x
455
+ if smallest is None:
456
+ return
457
+ if 'area' in smallest[1]:
458
+ if smallest[1]['area'] == c_area:
459
+ return
460
+ n = c[1].copy()
461
+ conds += [[smallest[0], n]]
462
+
463
+ def calculate_start_end_timesteps(model, conds):
464
+ for t in range(len(conds)):
465
+ x = conds[t]
466
+
467
+ timestep_start = None
468
+ timestep_end = None
469
+ if 'start_percent' in x[1]:
470
+ timestep_start = model.sigma_to_t(model.t_to_sigma(torch.tensor(x[1]['start_percent'] * 999.0)))
471
+ if 'end_percent' in x[1]:
472
+ timestep_end = model.sigma_to_t(model.t_to_sigma(torch.tensor(x[1]['end_percent'] * 999.0)))
473
+
474
+ if (timestep_start is not None) or (timestep_end is not None):
475
+ n = x[1].copy()
476
+ if (timestep_start is not None):
477
+ n['timestep_start'] = timestep_start
478
+ if (timestep_end is not None):
479
+ n['timestep_end'] = timestep_end
480
+ conds[t] = [x[0], n]
481
+
482
+ def pre_run_control(model, conds):
483
+ for t in range(len(conds)):
484
+ x = conds[t]
485
+
486
+ timestep_start = None
487
+ timestep_end = None
488
+ percent_to_timestep_function = lambda a: model.sigma_to_t(model.t_to_sigma(torch.tensor(a) * 999.0))
489
+ if 'control' in x[1]:
490
+ x[1]['control'].pre_run(model.inner_model.inner_model, percent_to_timestep_function)
491
+
492
+ def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
493
+ cond_cnets = []
494
+ cond_other = []
495
+ uncond_cnets = []
496
+ uncond_other = []
497
+ for t in range(len(conds)):
498
+ x = conds[t]
499
+ if 'area' not in x[1]:
500
+ if name in x[1] and x[1][name] is not None:
501
+ cond_cnets.append(x[1][name])
502
+ else:
503
+ cond_other.append((x, t))
504
+ for t in range(len(uncond)):
505
+ x = uncond[t]
506
+ if 'area' not in x[1]:
507
+ if name in x[1] and x[1][name] is not None:
508
+ uncond_cnets.append(x[1][name])
509
+ else:
510
+ uncond_other.append((x, t))
511
+
512
+ if len(uncond_cnets) > 0:
513
+ return
514
+
515
+ for x in range(len(cond_cnets)):
516
+ temp = uncond_other[x % len(uncond_other)]
517
+ o = temp[0]
518
+ if name in o[1] and o[1][name] is not None:
519
+ n = o[1].copy()
520
+ n[name] = uncond_fill_func(cond_cnets, x)
521
+ uncond += [[o[0], n]]
522
+ else:
523
+ n = o[1].copy()
524
+ n[name] = uncond_fill_func(cond_cnets, x)
525
+ uncond[temp[1]] = [o[0], n]
526
+
527
+ def encode_adm(model, conds, batch_size, width, height, device, prompt_type):
528
+ for t in range(len(conds)):
529
+ x = conds[t]
530
+ adm_out = None
531
+ if 'adm' in x[1]:
532
+ adm_out = x[1]["adm"]
533
+ else:
534
+ params = x[1].copy()
535
+ params["width"] = params.get("width", width * 8)
536
+ params["height"] = params.get("height", height * 8)
537
+ params["prompt_type"] = params.get("prompt_type", prompt_type)
538
+ adm_out = model.encode_adm(device=device, **params)
539
+
540
+ if adm_out is not None:
541
+ x[1] = x[1].copy()
542
+ x[1]["adm_encoded"] = comfy.utils.repeat_to_batch_size(adm_out, batch_size).to(device)
543
+
544
+ return conds
545
+
546
+
547
+ class KSampler:
548
+ SCHEDULERS = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
549
+ SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
550
+ "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
551
+ "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "ddim", "uni_pc", "uni_pc_bh2"]
552
+
553
+ def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
554
+ self.model = model
555
+ self.model_denoise = CFGNoisePredictor(self.model)
556
+ if self.model.model_type == model_base.ModelType.V_PREDICTION:
557
+ self.model_wrap = CompVisVDenoiser(self.model_denoise, quantize=True)
558
+ else:
559
+ self.model_wrap = k_diffusion_external.CompVisDenoiser(self.model_denoise, quantize=True)
560
+
561
+ self.model_k = KSamplerX0Inpaint(self.model_wrap)
562
+ self.device = device
563
+ if scheduler not in self.SCHEDULERS:
564
+ scheduler = self.SCHEDULERS[0]
565
+ if sampler not in self.SAMPLERS:
566
+ sampler = self.SAMPLERS[0]
567
+ self.scheduler = scheduler
568
+ self.sampler = sampler
569
+ self.sigma_min=float(self.model_wrap.sigma_min)
570
+ self.sigma_max=float(self.model_wrap.sigma_max)
571
+ self.set_steps(steps, denoise)
572
+ self.denoise = denoise
573
+ self.model_options = model_options
574
+
575
+ def calculate_sigmas(self, steps):
576
+ sigmas = None
577
+
578
+ discard_penultimate_sigma = False
579
+ if self.sampler in ['dpm_2', 'dpm_2_ancestral']:
580
+ steps += 1
581
+ discard_penultimate_sigma = True
582
+
583
+ if self.scheduler == "karras":
584
+ sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
585
+ elif self.scheduler == "exponential":
586
+ sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
587
+ elif self.scheduler == "normal":
588
+ sigmas = self.model_wrap.get_sigmas(steps)
589
+ elif self.scheduler == "simple":
590
+ sigmas = simple_scheduler(self.model_wrap, steps)
591
+ elif self.scheduler == "ddim_uniform":
592
+ sigmas = ddim_scheduler(self.model_wrap, steps)
593
+ elif self.scheduler == "sgm_uniform":
594
+ sigmas = sgm_scheduler(self.model_wrap, steps)
595
+ else:
596
+ print("error invalid scheduler", self.scheduler)
597
+
598
+ if discard_penultimate_sigma:
599
+ sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
600
+ return sigmas
601
+
602
+ def set_steps(self, steps, denoise=None):
603
+ self.steps = steps
604
+ if denoise is None or denoise > 0.9999:
605
+ self.sigmas = self.calculate_sigmas(steps).to(self.device)
606
+ else:
607
+ new_steps = int(steps/denoise)
608
+ sigmas = self.calculate_sigmas(new_steps).to(self.device)
609
+ self.sigmas = sigmas[-(steps + 1):]
610
+
611
+ def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
612
+ if sigmas is None:
613
+ sigmas = self.sigmas
614
+ sigma_min = self.sigma_min
615
+
616
+ if last_step is not None and last_step < (len(sigmas) - 1):
617
+ sigma_min = sigmas[last_step]
618
+ sigmas = sigmas[:last_step + 1]
619
+ if force_full_denoise:
620
+ sigmas[-1] = 0
621
+
622
+ if start_step is not None:
623
+ if start_step < (len(sigmas) - 1):
624
+ sigmas = sigmas[start_step:]
625
+ else:
626
+ if latent_image is not None:
627
+ return latent_image
628
+ else:
629
+ return torch.zeros_like(noise)
630
+
631
+ positive = positive[:]
632
+ negative = negative[:]
633
+
634
+ resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], self.device)
635
+ resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], self.device)
636
+
637
+ calculate_start_end_timesteps(self.model_wrap, negative)
638
+ calculate_start_end_timesteps(self.model_wrap, positive)
639
+
640
+ #make sure each cond area has an opposite one with the same area
641
+ for c in positive:
642
+ create_cond_with_same_area_if_none(negative, c)
643
+ for c in negative:
644
+ create_cond_with_same_area_if_none(positive, c)
645
+
646
+ pre_run_control(self.model_wrap, negative + positive)
647
+
648
+ apply_empty_x_to_equal_area(list(filter(lambda c: c[1].get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
649
+ apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
650
+
651
+ if self.model.is_adm():
652
+ positive = encode_adm(self.model, positive, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "positive")
653
+ negative = encode_adm(self.model, negative, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "negative")
654
+
655
+ if latent_image is not None:
656
+ latent_image = self.model.process_latent_in(latent_image)
657
+
658
+ extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": self.model_options, "seed":seed}
659
+
660
+ cond_concat = None
661
+ if hasattr(self.model, 'concat_keys'): #inpaint
662
+ cond_concat = []
663
+ for ck in self.model.concat_keys:
664
+ if denoise_mask is not None:
665
+ if ck == "mask":
666
+ cond_concat.append(denoise_mask[:,:1])
667
+ elif ck == "masked_image":
668
+ cond_concat.append(latent_image) #NOTE: the latent_image should be masked by the mask in pixel space
669
+ else:
670
+ if ck == "mask":
671
+ cond_concat.append(torch.ones_like(noise)[:,:1])
672
+ elif ck == "masked_image":
673
+ cond_concat.append(blank_inpaint_image_like(noise))
674
+ extra_args["cond_concat"] = cond_concat
675
+
676
+ if sigmas[0] != self.sigmas[0] or (self.denoise is not None and self.denoise < 1.0):
677
+ max_denoise = False
678
+ else:
679
+ max_denoise = True
680
+
681
+
682
+ if self.sampler == "uni_pc":
683
+ samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)
684
+ elif self.sampler == "uni_pc_bh2":
685
+ samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
686
+ elif self.sampler == "ddim":
687
+ timesteps = []
688
+ for s in range(sigmas.shape[0]):
689
+ timesteps.insert(0, self.model_wrap.sigma_to_discrete_timestep(sigmas[s]))
690
+ noise_mask = None
691
+ if denoise_mask is not None:
692
+ noise_mask = 1.0 - denoise_mask
693
+
694
+ ddim_callback = None
695
+ if callback is not None:
696
+ total_steps = len(timesteps) - 1
697
+ ddim_callback = lambda pred_x0, i: callback(i, pred_x0, None, total_steps)
698
+
699
+ sampler = DDIMSampler(self.model, device=self.device)
700
+ sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
701
+ z_enc = sampler.stochastic_encode(latent_image, torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(self.device), noise=noise, max_denoise=max_denoise)
702
+ samples, _ = sampler.sample_custom(ddim_timesteps=timesteps,
703
+ conditioning=positive,
704
+ batch_size=noise.shape[0],
705
+ shape=noise.shape[1:],
706
+ verbose=False,
707
+ unconditional_guidance_scale=cfg,
708
+ unconditional_conditioning=negative,
709
+ eta=0.0,
710
+ x_T=z_enc,
711
+ x0=latent_image,
712
+ img_callback=ddim_callback,
713
+ denoise_function=self.model_wrap.predict_eps_discrete_timestep,
714
+ extra_args=extra_args,
715
+ mask=noise_mask,
716
+ to_zero=sigmas[-1]==0,
717
+ end_step=sigmas.shape[0] - 1,
718
+ disable_pbar=disable_pbar)
719
+
720
+ else:
721
+ extra_args["denoise_mask"] = denoise_mask
722
+ self.model_k.latent_image = latent_image
723
+ self.model_k.noise = noise
724
+
725
+ if max_denoise:
726
+ noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
727
+ else:
728
+ noise = noise * sigmas[0]
729
+
730
+ k_callback = None
731
+ total_steps = len(sigmas) - 1
732
+ if callback is not None:
733
+ k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
734
+
735
+ if latent_image is not None:
736
+ noise += latent_image
737
+ if self.sampler == "dpm_fast":
738
+ samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
739
+ elif self.sampler == "dpm_adaptive":
740
+ samples = k_diffusion_sampling.sample_dpm_adaptive(self.model_k, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=k_callback, disable=disable_pbar)
741
+ else:
742
+ samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
743
+
744
+ return self.model.process_latent_out(samples.to(torch.float32))
comfy/sd.py ADDED
@@ -0,0 +1,486 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import contextlib
3
+ import math
4
+
5
+ from comfy import model_management
6
+ from .ldm.util import instantiate_from_config
7
+ from .ldm.models.autoencoder import AutoencoderKL
8
+ import yaml
9
+
10
+ import comfy.utils
11
+
12
+ from . import clip_vision
13
+ from . import gligen
14
+ from . import diffusers_convert
15
+ from . import model_base
16
+ from . import model_detection
17
+
18
+ from . import sd1_clip
19
+ from . import sd2_clip
20
+ from . import sdxl_clip
21
+
22
+ import comfy.model_patcher
23
+ import comfy.lora
24
+ import comfy.t2i_adapter.adapter
25
+ import comfy.supported_models_base
26
+
27
+ def load_model_weights(model, sd):
28
+ m, u = model.load_state_dict(sd, strict=False)
29
+ m = set(m)
30
+ unexpected_keys = set(u)
31
+
32
+ k = list(sd.keys())
33
+ for x in k:
34
+ if x not in unexpected_keys:
35
+ w = sd.pop(x)
36
+ del w
37
+ if len(m) > 0:
38
+ print("missing", m)
39
+ return model
40
+
41
+ def load_clip_weights(model, sd):
42
+ k = list(sd.keys())
43
+ for x in k:
44
+ if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):
45
+ y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")
46
+ sd[y] = sd.pop(x)
47
+
48
+ if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in sd:
49
+ ids = sd['cond_stage_model.transformer.text_model.embeddings.position_ids']
50
+ if ids.dtype == torch.float32:
51
+ sd['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
52
+
53
+ sd = comfy.utils.transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.text_model.", 24)
54
+ return load_model_weights(model, sd)
55
+
56
+
57
+ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
58
+ key_map = comfy.lora.model_lora_keys_unet(model.model)
59
+ key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
60
+ loaded = comfy.lora.load_lora(lora, key_map)
61
+ new_modelpatcher = model.clone()
62
+ k = new_modelpatcher.add_patches(loaded, strength_model)
63
+ new_clip = clip.clone()
64
+ k1 = new_clip.add_patches(loaded, strength_clip)
65
+ k = set(k)
66
+ k1 = set(k1)
67
+ for x in loaded:
68
+ if (x not in k) and (x not in k1):
69
+ print("NOT LOADED", x)
70
+
71
+ return (new_modelpatcher, new_clip)
72
+
73
+
74
+ class CLIP:
75
+ def __init__(self, target=None, embedding_directory=None, no_init=False):
76
+ if no_init:
77
+ return
78
+ params = target.params.copy()
79
+ clip = target.clip
80
+ tokenizer = target.tokenizer
81
+
82
+ load_device = model_management.text_encoder_device()
83
+ offload_device = model_management.text_encoder_offload_device()
84
+ params['device'] = offload_device
85
+ if model_management.should_use_fp16(load_device, prioritize_performance=False):
86
+ params['dtype'] = torch.float16
87
+ else:
88
+ params['dtype'] = torch.float32
89
+
90
+ self.cond_stage_model = clip(**(params))
91
+
92
+ self.tokenizer = tokenizer(embedding_directory=embedding_directory)
93
+ self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
94
+ self.layer_idx = None
95
+
96
+ def clone(self):
97
+ n = CLIP(no_init=True)
98
+ n.patcher = self.patcher.clone()
99
+ n.cond_stage_model = self.cond_stage_model
100
+ n.tokenizer = self.tokenizer
101
+ n.layer_idx = self.layer_idx
102
+ return n
103
+
104
+ def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
105
+ return self.patcher.add_patches(patches, strength_patch, strength_model)
106
+
107
+ def clip_layer(self, layer_idx):
108
+ self.layer_idx = layer_idx
109
+
110
+ def tokenize(self, text, return_word_ids=False):
111
+ return self.tokenizer.tokenize_with_weights(text, return_word_ids)
112
+
113
+ def encode_from_tokens(self, tokens, return_pooled=False):
114
+ if self.layer_idx is not None:
115
+ self.cond_stage_model.clip_layer(self.layer_idx)
116
+ else:
117
+ self.cond_stage_model.reset_clip_layer()
118
+
119
+ self.load_model()
120
+ cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
121
+ if return_pooled:
122
+ return cond, pooled
123
+ return cond
124
+
125
+ def encode(self, text):
126
+ tokens = self.tokenize(text)
127
+ return self.encode_from_tokens(tokens)
128
+
129
+ def load_sd(self, sd):
130
+ return self.cond_stage_model.load_sd(sd)
131
+
132
+ def get_sd(self):
133
+ return self.cond_stage_model.state_dict()
134
+
135
+ def load_model(self):
136
+ model_management.load_model_gpu(self.patcher)
137
+ return self.patcher
138
+
139
+ def get_key_patches(self):
140
+ return self.patcher.get_key_patches()
141
+
142
+ class VAE:
143
+ def __init__(self, ckpt_path=None, device=None, config=None):
144
+ if config is None:
145
+ #default SD1.x/SD2.x VAE parameters
146
+ ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
147
+ self.first_stage_model = AutoencoderKL(ddconfig, {'target': 'torch.nn.Identity'}, 4, monitor="val/rec_loss")
148
+ else:
149
+ self.first_stage_model = AutoencoderKL(**(config['params']))
150
+ self.first_stage_model = self.first_stage_model.eval()
151
+ if ckpt_path is not None:
152
+ sd = comfy.utils.load_torch_file(ckpt_path)
153
+ if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
154
+ sd = diffusers_convert.convert_vae_state_dict(sd)
155
+ self.first_stage_model.load_state_dict(sd, strict=False)
156
+
157
+ if device is None:
158
+ device = model_management.vae_device()
159
+ self.device = device
160
+ self.offload_device = model_management.vae_offload_device()
161
+ self.vae_dtype = model_management.vae_dtype()
162
+ self.first_stage_model.to(self.vae_dtype)
163
+
164
+ def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
165
+ steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
166
+ steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
167
+ steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
168
+ pbar = comfy.utils.ProgressBar(steps)
169
+
170
+ decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
171
+ output = torch.clamp((
172
+ (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = 8, pbar = pbar) +
173
+ comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = 8, pbar = pbar) +
174
+ comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = 8, pbar = pbar))
175
+ / 3.0) / 2.0, min=0.0, max=1.0)
176
+ return output
177
+
178
+ def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
179
+ steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
180
+ steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
181
+ steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)
182
+ pbar = comfy.utils.ProgressBar(steps)
183
+
184
+ encode_fn = lambda a: self.first_stage_model.encode(2. * a.to(self.vae_dtype).to(self.device) - 1.).sample().float()
185
+ samples = comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
186
+ samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
187
+ samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
188
+ samples /= 3.0
189
+ return samples
190
+
191
+ def decode(self, samples_in):
192
+ self.first_stage_model = self.first_stage_model.to(self.device)
193
+ try:
194
+ memory_used = (2562 * samples_in.shape[2] * samples_in.shape[3] * 64) * 1.7
195
+ model_management.free_memory(memory_used, self.device)
196
+ free_memory = model_management.get_free_memory(self.device)
197
+ batch_number = int(free_memory / memory_used)
198
+ batch_number = max(1, batch_number)
199
+
200
+ pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * 8), round(samples_in.shape[3] * 8)), device="cpu")
201
+ for x in range(0, samples_in.shape[0], batch_number):
202
+ samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
203
+ pixel_samples[x:x+batch_number] = torch.clamp((self.first_stage_model.decode(samples) + 1.0) / 2.0, min=0.0, max=1.0).cpu().float()
204
+ except model_management.OOM_EXCEPTION as e:
205
+ print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
206
+ pixel_samples = self.decode_tiled_(samples_in)
207
+
208
+ self.first_stage_model = self.first_stage_model.to(self.offload_device)
209
+ pixel_samples = pixel_samples.cpu().movedim(1,-1)
210
+ return pixel_samples
211
+
212
+ def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
213
+ self.first_stage_model = self.first_stage_model.to(self.device)
214
+ output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
215
+ self.first_stage_model = self.first_stage_model.to(self.offload_device)
216
+ return output.movedim(1,-1)
217
+
218
+ def encode(self, pixel_samples):
219
+ self.first_stage_model = self.first_stage_model.to(self.device)
220
+ pixel_samples = pixel_samples.movedim(-1,1)
221
+ try:
222
+ memory_used = (2078 * pixel_samples.shape[2] * pixel_samples.shape[3]) * 1.7 #NOTE: this constant along with the one in the decode above are estimated from the mem usage for the VAE and could change.
223
+ model_management.free_memory(memory_used, self.device)
224
+ free_memory = model_management.get_free_memory(self.device)
225
+ batch_number = int(free_memory / memory_used)
226
+ batch_number = max(1, batch_number)
227
+ samples = torch.empty((pixel_samples.shape[0], 4, round(pixel_samples.shape[2] // 8), round(pixel_samples.shape[3] // 8)), device="cpu")
228
+ for x in range(0, pixel_samples.shape[0], batch_number):
229
+ pixels_in = (2. * pixel_samples[x:x+batch_number] - 1.).to(self.vae_dtype).to(self.device)
230
+ samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).sample().cpu().float()
231
+
232
+ except model_management.OOM_EXCEPTION as e:
233
+ print("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
234
+ samples = self.encode_tiled_(pixel_samples)
235
+
236
+ self.first_stage_model = self.first_stage_model.to(self.offload_device)
237
+ return samples
238
+
239
+ def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
240
+ self.first_stage_model = self.first_stage_model.to(self.device)
241
+ pixel_samples = pixel_samples.movedim(-1,1)
242
+ samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
243
+ self.first_stage_model = self.first_stage_model.to(self.offload_device)
244
+ return samples
245
+
246
+ def get_sd(self):
247
+ return self.first_stage_model.state_dict()
248
+
249
+ class StyleModel:
250
+ def __init__(self, model, device="cpu"):
251
+ self.model = model
252
+
253
+ def get_cond(self, input):
254
+ return self.model(input.last_hidden_state)
255
+
256
+
257
+ def load_style_model(ckpt_path):
258
+ model_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
259
+ keys = model_data.keys()
260
+ if "style_embedding" in keys:
261
+ model = comfy.t2i_adapter.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
262
+ else:
263
+ raise Exception("invalid style model {}".format(ckpt_path))
264
+ model.load_state_dict(model_data)
265
+ return StyleModel(model)
266
+
267
+
268
+ def load_clip(ckpt_paths, embedding_directory=None):
269
+ clip_data = []
270
+ for p in ckpt_paths:
271
+ clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
272
+
273
+ class EmptyClass:
274
+ pass
275
+
276
+ for i in range(len(clip_data)):
277
+ if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
278
+ clip_data[i] = comfy.utils.transformers_convert(clip_data[i], "", "text_model.", 32)
279
+
280
+ clip_target = EmptyClass()
281
+ clip_target.params = {}
282
+ if len(clip_data) == 1:
283
+ if "text_model.encoder.layers.30.mlp.fc1.weight" in clip_data[0]:
284
+ clip_target.clip = sdxl_clip.SDXLRefinerClipModel
285
+ clip_target.tokenizer = sdxl_clip.SDXLTokenizer
286
+ elif "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data[0]:
287
+ clip_target.clip = sd2_clip.SD2ClipModel
288
+ clip_target.tokenizer = sd2_clip.SD2Tokenizer
289
+ else:
290
+ clip_target.clip = sd1_clip.SD1ClipModel
291
+ clip_target.tokenizer = sd1_clip.SD1Tokenizer
292
+ else:
293
+ clip_target.clip = sdxl_clip.SDXLClipModel
294
+ clip_target.tokenizer = sdxl_clip.SDXLTokenizer
295
+
296
+ clip = CLIP(clip_target, embedding_directory=embedding_directory)
297
+ for c in clip_data:
298
+ m, u = clip.load_sd(c)
299
+ if len(m) > 0:
300
+ print("clip missing:", m)
301
+
302
+ if len(u) > 0:
303
+ print("clip unexpected:", u)
304
+ return clip
305
+
306
+ def load_gligen(ckpt_path):
307
+ data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
308
+ model = gligen.load_gligen(data)
309
+ if model_management.should_use_fp16():
310
+ model = model.half()
311
+ return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
312
+
313
+ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
314
+ #TODO: this function is a mess and should be removed eventually
315
+ if config is None:
316
+ with open(config_path, 'r') as stream:
317
+ config = yaml.safe_load(stream)
318
+ model_config_params = config['model']['params']
319
+ clip_config = model_config_params['cond_stage_config']
320
+ scale_factor = model_config_params['scale_factor']
321
+ vae_config = model_config_params['first_stage_config']
322
+
323
+ fp16 = False
324
+ if "unet_config" in model_config_params:
325
+ if "params" in model_config_params["unet_config"]:
326
+ unet_config = model_config_params["unet_config"]["params"]
327
+ if "use_fp16" in unet_config:
328
+ fp16 = unet_config["use_fp16"]
329
+
330
+ noise_aug_config = None
331
+ if "noise_aug_config" in model_config_params:
332
+ noise_aug_config = model_config_params["noise_aug_config"]
333
+
334
+ model_type = model_base.ModelType.EPS
335
+
336
+ if "parameterization" in model_config_params:
337
+ if model_config_params["parameterization"] == "v":
338
+ model_type = model_base.ModelType.V_PREDICTION
339
+
340
+ clip = None
341
+ vae = None
342
+
343
+ class WeightsLoader(torch.nn.Module):
344
+ pass
345
+
346
+ if state_dict is None:
347
+ state_dict = comfy.utils.load_torch_file(ckpt_path)
348
+
349
+ class EmptyClass:
350
+ pass
351
+
352
+ model_config = comfy.supported_models_base.BASE({})
353
+
354
+ from . import latent_formats
355
+ model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)
356
+ model_config.unet_config = unet_config
357
+
358
+ if config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
359
+ model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
360
+ else:
361
+ model = model_base.BaseModel(model_config, model_type=model_type)
362
+
363
+ if config['model']["target"].endswith("LatentInpaintDiffusion"):
364
+ model.set_inpaint()
365
+
366
+ if fp16:
367
+ model = model.half()
368
+
369
+ offload_device = model_management.unet_offload_device()
370
+ model = model.to(offload_device)
371
+ model.load_model_weights(state_dict, "model.diffusion_model.")
372
+
373
+ if output_vae:
374
+ w = WeightsLoader()
375
+ vae = VAE(config=vae_config)
376
+ w.first_stage_model = vae.first_stage_model
377
+ load_model_weights(w, state_dict)
378
+
379
+ if output_clip:
380
+ w = WeightsLoader()
381
+ clip_target = EmptyClass()
382
+ clip_target.params = clip_config.get("params", {})
383
+ if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
384
+ clip_target.clip = sd2_clip.SD2ClipModel
385
+ clip_target.tokenizer = sd2_clip.SD2Tokenizer
386
+ elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
387
+ clip_target.clip = sd1_clip.SD1ClipModel
388
+ clip_target.tokenizer = sd1_clip.SD1Tokenizer
389
+ clip = CLIP(clip_target, embedding_directory=embedding_directory)
390
+ w.cond_stage_model = clip.cond_stage_model
391
+ load_clip_weights(w, state_dict)
392
+
393
+ return (comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
394
+
395
+ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None):
396
+ sd = comfy.utils.load_torch_file(ckpt_path)
397
+ sd_keys = sd.keys()
398
+ clip = None
399
+ clipvision = None
400
+ vae = None
401
+ model = None
402
+ clip_target = None
403
+
404
+ parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.")
405
+ fp16 = model_management.should_use_fp16(model_params=parameters)
406
+
407
+ class WeightsLoader(torch.nn.Module):
408
+ pass
409
+
410
+ model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", fp16)
411
+ if model_config is None:
412
+ raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
413
+
414
+ if model_config.clip_vision_prefix is not None:
415
+ if output_clipvision:
416
+ clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
417
+
418
+ dtype = torch.float32
419
+ if fp16:
420
+ dtype = torch.float16
421
+
422
+ inital_load_device = model_management.unet_inital_load_device(parameters, dtype)
423
+ offload_device = model_management.unet_offload_device()
424
+ model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
425
+ model.load_model_weights(sd, "model.diffusion_model.")
426
+
427
+ if output_vae:
428
+ vae = VAE()
429
+ w = WeightsLoader()
430
+ w.first_stage_model = vae.first_stage_model
431
+ load_model_weights(w, sd)
432
+
433
+ if output_clip:
434
+ w = WeightsLoader()
435
+ clip_target = model_config.clip_target()
436
+ clip = CLIP(clip_target, embedding_directory=embedding_directory)
437
+ w.cond_stage_model = clip.cond_stage_model
438
+ sd = model_config.process_clip_state_dict(sd)
439
+ load_model_weights(w, sd)
440
+
441
+ left_over = sd.keys()
442
+ if len(left_over) > 0:
443
+ print("left over keys:", left_over)
444
+
445
+ model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
446
+ if inital_load_device != torch.device("cpu"):
447
+ print("loaded straight to GPU")
448
+ model_management.load_model_gpu(model_patcher)
449
+
450
+ return (model_patcher, clip, vae, clipvision)
451
+
452
+
453
+ def load_unet(unet_path): #load unet in diffusers format
454
+ sd = comfy.utils.load_torch_file(unet_path)
455
+ parameters = comfy.utils.calculate_parameters(sd)
456
+ fp16 = model_management.should_use_fp16(model_params=parameters)
457
+ if "input_blocks.0.0.weight" in sd: #ldm
458
+ model_config = model_detection.model_config_from_unet(sd, "", fp16)
459
+ if model_config is None:
460
+ raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
461
+ new_sd = sd
462
+
463
+ else: #diffusers
464
+ model_config = model_detection.model_config_from_diffusers_unet(sd, fp16)
465
+ if model_config is None:
466
+ print("ERROR UNSUPPORTED UNET", unet_path)
467
+ return None
468
+
469
+ diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)
470
+
471
+ new_sd = {}
472
+ for k in diffusers_keys:
473
+ if k in sd:
474
+ new_sd[diffusers_keys[k]] = sd.pop(k)
475
+ else:
476
+ print(diffusers_keys[k], k)
477
+ offload_device = model_management.unet_offload_device()
478
+ model = model_config.get_model(new_sd, "")
479
+ model = model.to(offload_device)
480
+ model.load_model_weights(new_sd, "")
481
+ return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device)
482
+
483
+ def save_checkpoint(output_path, model, clip, vae, metadata=None):
484
+ model_management.load_models_gpu([model, clip.load_model()])
485
+ sd = model.model.state_dict_for_saving(clip.get_sd(), vae.get_sd())
486
+ comfy.utils.save_torch_file(sd, output_path, metadata=metadata)
comfy/sd1_clip.py ADDED
@@ -0,0 +1,450 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextConfig, modeling_utils
4
+ import comfy.ops
5
+ import torch
6
+ import traceback
7
+ import zipfile
8
+ from . import model_management
9
+ import contextlib
10
+
11
+ class ClipTokenWeightEncoder:
12
+ def encode_token_weights(self, token_weight_pairs):
13
+ to_encode = list(self.empty_tokens)
14
+ for x in token_weight_pairs:
15
+ tokens = list(map(lambda a: a[0], x))
16
+ to_encode.append(tokens)
17
+
18
+ out, pooled = self.encode(to_encode)
19
+ z_empty = out[0:1]
20
+ if pooled.shape[0] > 1:
21
+ first_pooled = pooled[1:2]
22
+ else:
23
+ first_pooled = pooled[0:1]
24
+
25
+ output = []
26
+ for k in range(1, out.shape[0]):
27
+ z = out[k:k+1]
28
+ for i in range(len(z)):
29
+ for j in range(len(z[i])):
30
+ weight = token_weight_pairs[k - 1][j][1]
31
+ z[i][j] = (z[i][j] - z_empty[0][j]) * weight + z_empty[0][j]
32
+ output.append(z)
33
+
34
+ if (len(output) == 0):
35
+ return z_empty.cpu(), first_pooled.cpu()
36
+ return torch.cat(output, dim=-2).cpu(), first_pooled.cpu()
37
+
38
+ class SD1ClipModel(torch.nn.Module, ClipTokenWeightEncoder):
39
+ """Uses the CLIP transformer encoder for text (from huggingface)"""
40
+ LAYERS = [
41
+ "last",
42
+ "pooled",
43
+ "hidden"
44
+ ]
45
+ def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
46
+ freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, textmodel_path=None, dtype=None): # clip-vit-base-patch32
47
+ super().__init__()
48
+ assert layer in self.LAYERS
49
+ self.num_layers = 12
50
+ if textmodel_path is not None:
51
+ self.transformer = CLIPTextModel.from_pretrained(textmodel_path)
52
+ else:
53
+ if textmodel_json_config is None:
54
+ textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")
55
+ config = CLIPTextConfig.from_json_file(textmodel_json_config)
56
+ self.num_layers = config.num_hidden_layers
57
+ with comfy.ops.use_comfy_ops(device, dtype):
58
+ with modeling_utils.no_init_weights():
59
+ self.transformer = CLIPTextModel(config)
60
+
61
+ if dtype is not None:
62
+ self.transformer.to(dtype)
63
+ self.transformer.text_model.embeddings.token_embedding.to(torch.float32)
64
+ self.transformer.text_model.embeddings.position_embedding.to(torch.float32)
65
+
66
+ self.max_length = max_length
67
+ if freeze:
68
+ self.freeze()
69
+ self.layer = layer
70
+ self.layer_idx = None
71
+ self.empty_tokens = [[49406] + [49407] * 76]
72
+ self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
73
+ self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
74
+ self.enable_attention_masks = False
75
+
76
+ self.layer_norm_hidden_state = True
77
+ if layer == "hidden":
78
+ assert layer_idx is not None
79
+ assert abs(layer_idx) <= self.num_layers
80
+ self.clip_layer(layer_idx)
81
+ self.layer_default = (self.layer, self.layer_idx)
82
+
83
+ def freeze(self):
84
+ self.transformer = self.transformer.eval()
85
+ #self.train = disabled_train
86
+ for param in self.parameters():
87
+ param.requires_grad = False
88
+
89
+ def clip_layer(self, layer_idx):
90
+ if abs(layer_idx) >= self.num_layers:
91
+ self.layer = "last"
92
+ else:
93
+ self.layer = "hidden"
94
+ self.layer_idx = layer_idx
95
+
96
+ def reset_clip_layer(self):
97
+ self.layer = self.layer_default[0]
98
+ self.layer_idx = self.layer_default[1]
99
+
100
+ def set_up_textual_embeddings(self, tokens, current_embeds):
101
+ out_tokens = []
102
+ next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
103
+ embedding_weights = []
104
+
105
+ for x in tokens:
106
+ tokens_temp = []
107
+ for y in x:
108
+ if isinstance(y, int):
109
+ if y == token_dict_size: #EOS token
110
+ y = -1
111
+ tokens_temp += [y]
112
+ else:
113
+ if y.shape[0] == current_embeds.weight.shape[1]:
114
+ embedding_weights += [y]
115
+ tokens_temp += [next_new_token]
116
+ next_new_token += 1
117
+ else:
118
+ print("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored", y.shape[0], current_embeds.weight.shape[1])
119
+ while len(tokens_temp) < len(x):
120
+ tokens_temp += [self.empty_tokens[0][-1]]
121
+ out_tokens += [tokens_temp]
122
+
123
+ n = token_dict_size
124
+ if len(embedding_weights) > 0:
125
+ new_embedding = torch.nn.Embedding(next_new_token + 1, current_embeds.weight.shape[1], device=current_embeds.weight.device, dtype=current_embeds.weight.dtype)
126
+ new_embedding.weight[:token_dict_size] = current_embeds.weight[:-1]
127
+ for x in embedding_weights:
128
+ new_embedding.weight[n] = x
129
+ n += 1
130
+ new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
131
+ self.transformer.set_input_embeddings(new_embedding)
132
+
133
+ processed_tokens = []
134
+ for x in out_tokens:
135
+ processed_tokens += [list(map(lambda a: n if a == -1 else a, x))] #The EOS token should always be the largest one
136
+
137
+ return processed_tokens
138
+
139
+ def forward(self, tokens):
140
+ backup_embeds = self.transformer.get_input_embeddings()
141
+ device = backup_embeds.weight.device
142
+ tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
143
+ tokens = torch.LongTensor(tokens).to(device)
144
+
145
+ if self.transformer.text_model.final_layer_norm.weight.dtype != torch.float32:
146
+ precision_scope = torch.autocast
147
+ else:
148
+ precision_scope = lambda a, b: contextlib.nullcontext(a)
149
+
150
+ with precision_scope(model_management.get_autocast_device(device), torch.float32):
151
+ attention_mask = None
152
+ if self.enable_attention_masks:
153
+ attention_mask = torch.zeros_like(tokens)
154
+ max_token = self.transformer.get_input_embeddings().weight.shape[0] - 1
155
+ for x in range(attention_mask.shape[0]):
156
+ for y in range(attention_mask.shape[1]):
157
+ attention_mask[x, y] = 1
158
+ if tokens[x, y] == max_token:
159
+ break
160
+
161
+ outputs = self.transformer(input_ids=tokens, attention_mask=attention_mask, output_hidden_states=self.layer=="hidden")
162
+ self.transformer.set_input_embeddings(backup_embeds)
163
+
164
+ if self.layer == "last":
165
+ z = outputs.last_hidden_state
166
+ elif self.layer == "pooled":
167
+ z = outputs.pooler_output[:, None, :]
168
+ else:
169
+ z = outputs.hidden_states[self.layer_idx]
170
+ if self.layer_norm_hidden_state:
171
+ z = self.transformer.text_model.final_layer_norm(z)
172
+
173
+ pooled_output = outputs.pooler_output
174
+ if self.text_projection is not None:
175
+ pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
176
+ return z.float(), pooled_output.float()
177
+
178
+ def encode(self, tokens):
179
+ return self(tokens)
180
+
181
+ def load_sd(self, sd):
182
+ if "text_projection" in sd:
183
+ self.text_projection[:] = sd.pop("text_projection")
184
+ if "text_projection.weight" in sd:
185
+ self.text_projection[:] = sd.pop("text_projection.weight").transpose(0, 1)
186
+ return self.transformer.load_state_dict(sd, strict=False)
187
+
188
+ def parse_parentheses(string):
189
+ result = []
190
+ current_item = ""
191
+ nesting_level = 0
192
+ for char in string:
193
+ if char == "(":
194
+ if nesting_level == 0:
195
+ if current_item:
196
+ result.append(current_item)
197
+ current_item = "("
198
+ else:
199
+ current_item = "("
200
+ else:
201
+ current_item += char
202
+ nesting_level += 1
203
+ elif char == ")":
204
+ nesting_level -= 1
205
+ if nesting_level == 0:
206
+ result.append(current_item + ")")
207
+ current_item = ""
208
+ else:
209
+ current_item += char
210
+ else:
211
+ current_item += char
212
+ if current_item:
213
+ result.append(current_item)
214
+ return result
215
+
216
+ def token_weights(string, current_weight):
217
+ a = parse_parentheses(string)
218
+ out = []
219
+ for x in a:
220
+ weight = current_weight
221
+ if len(x) >= 2 and x[-1] == ')' and x[0] == '(':
222
+ x = x[1:-1]
223
+ xx = x.rfind(":")
224
+ weight *= 1.1
225
+ if xx > 0:
226
+ try:
227
+ weight = float(x[xx+1:])
228
+ x = x[:xx]
229
+ except:
230
+ pass
231
+ out += token_weights(x, weight)
232
+ else:
233
+ out += [(x, current_weight)]
234
+ return out
235
+
236
+ def escape_important(text):
237
+ text = text.replace("\\)", "\0\1")
238
+ text = text.replace("\\(", "\0\2")
239
+ return text
240
+
241
+ def unescape_important(text):
242
+ text = text.replace("\0\1", ")")
243
+ text = text.replace("\0\2", "(")
244
+ return text
245
+
246
+ def safe_load_embed_zip(embed_path):
247
+ with zipfile.ZipFile(embed_path) as myzip:
248
+ names = list(filter(lambda a: "data/" in a, myzip.namelist()))
249
+ names.reverse()
250
+ for n in names:
251
+ with myzip.open(n) as myfile:
252
+ data = myfile.read()
253
+ number = len(data) // 4
254
+ length_embed = 1024 #sd2.x
255
+ if number < 768:
256
+ continue
257
+ if number % 768 == 0:
258
+ length_embed = 768 #sd1.x
259
+ num_embeds = number // length_embed
260
+ embed = torch.frombuffer(data, dtype=torch.float)
261
+ out = embed.reshape((num_embeds, length_embed)).clone()
262
+ del embed
263
+ return out
264
+
265
+ def expand_directory_list(directories):
266
+ dirs = set()
267
+ for x in directories:
268
+ dirs.add(x)
269
+ for root, subdir, file in os.walk(x, followlinks=True):
270
+ dirs.add(root)
271
+ return list(dirs)
272
+
273
+ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
274
+ if isinstance(embedding_directory, str):
275
+ embedding_directory = [embedding_directory]
276
+
277
+ embedding_directory = expand_directory_list(embedding_directory)
278
+
279
+ valid_file = None
280
+ for embed_dir in embedding_directory:
281
+ embed_path = os.path.join(embed_dir, embedding_name)
282
+ if not os.path.isfile(embed_path):
283
+ extensions = ['.safetensors', '.pt', '.bin']
284
+ for x in extensions:
285
+ t = embed_path + x
286
+ if os.path.isfile(t):
287
+ valid_file = t
288
+ break
289
+ else:
290
+ valid_file = embed_path
291
+ if valid_file is not None:
292
+ break
293
+
294
+ if valid_file is None:
295
+ return None
296
+
297
+ embed_path = valid_file
298
+
299
+ embed_out = None
300
+
301
+ try:
302
+ if embed_path.lower().endswith(".safetensors"):
303
+ import safetensors.torch
304
+ embed = safetensors.torch.load_file(embed_path, device="cpu")
305
+ else:
306
+ if 'weights_only' in torch.load.__code__.co_varnames:
307
+ try:
308
+ embed = torch.load(embed_path, weights_only=True, map_location="cpu")
309
+ except:
310
+ embed_out = safe_load_embed_zip(embed_path)
311
+ else:
312
+ embed = torch.load(embed_path, map_location="cpu")
313
+ except Exception as e:
314
+ print(traceback.format_exc())
315
+ print()
316
+ print("error loading embedding, skipping loading:", embedding_name)
317
+ return None
318
+
319
+ if embed_out is None:
320
+ if 'string_to_param' in embed:
321
+ values = embed['string_to_param'].values()
322
+ embed_out = next(iter(values))
323
+ elif isinstance(embed, list):
324
+ out_list = []
325
+ for x in range(len(embed)):
326
+ for k in embed[x]:
327
+ t = embed[x][k]
328
+ if t.shape[-1] != embedding_size:
329
+ continue
330
+ out_list.append(t.reshape(-1, t.shape[-1]))
331
+ embed_out = torch.cat(out_list, dim=0)
332
+ elif embed_key is not None and embed_key in embed:
333
+ embed_out = embed[embed_key]
334
+ else:
335
+ values = embed.values()
336
+ embed_out = next(iter(values))
337
+ return embed_out
338
+
339
+ class SD1Tokenizer:
340
+ def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l'):
341
+ if tokenizer_path is None:
342
+ tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
343
+ self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path)
344
+ self.max_length = max_length
345
+ self.max_tokens_per_section = self.max_length - 2
346
+
347
+ empty = self.tokenizer('')["input_ids"]
348
+ self.start_token = empty[0]
349
+ self.end_token = empty[1]
350
+ self.pad_with_end = pad_with_end
351
+ vocab = self.tokenizer.get_vocab()
352
+ self.inv_vocab = {v: k for k, v in vocab.items()}
353
+ self.embedding_directory = embedding_directory
354
+ self.max_word_length = 8
355
+ self.embedding_identifier = "embedding:"
356
+ self.embedding_size = embedding_size
357
+ self.embedding_key = embedding_key
358
+
359
+ def _try_get_embedding(self, embedding_name:str):
360
+ '''
361
+ Takes a potential embedding name and tries to retrieve it.
362
+ Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
363
+ '''
364
+ embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
365
+ if embed is None:
366
+ stripped = embedding_name.strip(',')
367
+ if len(stripped) < len(embedding_name):
368
+ embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
369
+ return (embed, embedding_name[len(stripped):])
370
+ return (embed, "")
371
+
372
+
373
+ def tokenize_with_weights(self, text:str, return_word_ids=False):
374
+ '''
375
+ Takes a prompt and converts it to a list of (token, weight, word id) elements.
376
+ Tokens can both be integer tokens and pre computed CLIP tensors.
377
+ Word id values are unique per word and embedding, where the id 0 is reserved for non word tokens.
378
+ Returned list has the dimensions NxM where M is the input size of CLIP
379
+ '''
380
+ if self.pad_with_end:
381
+ pad_token = self.end_token
382
+ else:
383
+ pad_token = 0
384
+
385
+ text = escape_important(text)
386
+ parsed_weights = token_weights(text, 1.0)
387
+
388
+ #tokenize words
389
+ tokens = []
390
+ for weighted_segment, weight in parsed_weights:
391
+ to_tokenize = unescape_important(weighted_segment).replace("\n", " ").split(' ')
392
+ to_tokenize = [x for x in to_tokenize if x != ""]
393
+ for word in to_tokenize:
394
+ #if we find an embedding, deal with the embedding
395
+ if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
396
+ embedding_name = word[len(self.embedding_identifier):].strip('\n')
397
+ embed, leftover = self._try_get_embedding(embedding_name)
398
+ if embed is None:
399
+ print(f"warning, embedding:{embedding_name} does not exist, ignoring")
400
+ else:
401
+ if len(embed.shape) == 1:
402
+ tokens.append([(embed, weight)])
403
+ else:
404
+ tokens.append([(embed[x], weight) for x in range(embed.shape[0])])
405
+ #if we accidentally have leftover text, continue parsing using leftover, else move on to next word
406
+ if leftover != "":
407
+ word = leftover
408
+ else:
409
+ continue
410
+ #parse word
411
+ tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][1:-1]])
412
+
413
+ #reshape token array to CLIP input size
414
+ batched_tokens = []
415
+ batch = [(self.start_token, 1.0, 0)]
416
+ batched_tokens.append(batch)
417
+ for i, t_group in enumerate(tokens):
418
+ #determine if we're going to try and keep the tokens in a single batch
419
+ is_large = len(t_group) >= self.max_word_length
420
+
421
+ while len(t_group) > 0:
422
+ if len(t_group) + len(batch) > self.max_length - 1:
423
+ remaining_length = self.max_length - len(batch) - 1
424
+ #break word in two and add end token
425
+ if is_large:
426
+ batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
427
+ batch.append((self.end_token, 1.0, 0))
428
+ t_group = t_group[remaining_length:]
429
+ #add end token and pad
430
+ else:
431
+ batch.append((self.end_token, 1.0, 0))
432
+ batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
433
+ #start new batch
434
+ batch = [(self.start_token, 1.0, 0)]
435
+ batched_tokens.append(batch)
436
+ else:
437
+ batch.extend([(t,w,i+1) for t,w in t_group])
438
+ t_group = []
439
+
440
+ #fill last batch
441
+ batch.extend([(self.end_token, 1.0, 0)] + [(pad_token, 1.0, 0)] * (self.max_length - len(batch) - 1))
442
+
443
+ if not return_word_ids:
444
+ batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
445
+
446
+ return batched_tokens
447
+
448
+
449
+ def untokenize(self, token_weight_pair):
450
+ return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
comfy/sd1_clip_config.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "openai/clip-vit-large-patch14",
3
+ "architectures": [
4
+ "CLIPTextModel"
5
+ ],
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": 0,
8
+ "dropout": 0.0,
9
+ "eos_token_id": 2,
10
+ "hidden_act": "quick_gelu",
11
+ "hidden_size": 768,
12
+ "initializer_factor": 1.0,
13
+ "initializer_range": 0.02,
14
+ "intermediate_size": 3072,
15
+ "layer_norm_eps": 1e-05,
16
+ "max_position_embeddings": 77,
17
+ "model_type": "clip_text_model",
18
+ "num_attention_heads": 12,
19
+ "num_hidden_layers": 12,
20
+ "pad_token_id": 1,
21
+ "projection_dim": 768,
22
+ "torch_dtype": "float32",
23
+ "transformers_version": "4.24.0",
24
+ "vocab_size": 49408
25
+ }