samplers.py 32.3 KB
Newer Older
1
2
from .k_diffusion import sampling as k_diffusion_sampling
from .k_diffusion import external as k_diffusion_external
3
from .extra_samplers import uni_pc
comfyanonymous's avatar
comfyanonymous committed
4
import torch
5
from comfy import model_management
comfyanonymous's avatar
comfyanonymous committed
6
7
from .ldm.models.diffusion.ddim import DDIMSampler
from .ldm.modules.diffusionmodules.util import make_ddim_timesteps
comfyanonymous's avatar
comfyanonymous committed
8
import math
9
from comfy import model_base
10
import comfy.utils
comfyanonymous's avatar
comfyanonymous committed
11
12
13

def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9)
    return abs(a*b) // math.gcd(a, b)
comfyanonymous's avatar
comfyanonymous committed
14

comfyanonymous's avatar
comfyanonymous committed
15
16
#The main sampling function shared by all the samplers
#Returns predicted noise
17
18
def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
        def get_area_and_mult(cond, x_in, timestep_in):
19
20
            area = (x_in.shape[2], x_in.shape[3], 0, 0)
            strength = 1.0
21
22
            if 'timestep_start' in cond[1]:
                timestep_start = cond[1]['timestep_start']
23
                if timestep_in[0] > timestep_start:
24
25
26
                    return None
            if 'timestep_end' in cond[1]:
                timestep_end = cond[1]['timestep_end']
27
                if timestep_in[0] < timestep_end:
28
                    return None
29
30
31
32
            if 'area' in cond[1]:
                area = cond[1]['area']
            if 'strength' in cond[1]:
                strength = cond[1]['strength']
33

34
            adm_cond = None
35
36
            if 'adm_encoded' in cond[1]:
                adm_cond = cond[1]['adm_encoded']
37

38
            input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
Jacob Segal's avatar
Jacob Segal committed
39
40
41
            if 'mask' in cond[1]:
                # Scale the mask to the size of the input
                # The mask should have been resized as we began the sampling process
42
43
44
                mask_strength = 1.0
                if "mask_strength" in cond[1]:
                    mask_strength = cond[1]["mask_strength"]
Jacob Segal's avatar
Jacob Segal committed
45
46
47
                mask = cond[1]['mask']
                assert(mask.shape[1] == x_in.shape[2])
                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
Jacob Segal's avatar
Jacob Segal committed
49
                mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
Jacob Segal's avatar
Jacob Segal committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
            else:
                mask = torch.ones_like(input_x)
            mult = mask * strength

            if 'mask' not in cond[1]:
                rr = 8
                if area[2] != 0:
                    for t in range(rr):
                        mult[:,:,t:1+t,:] *= ((1.0/rr) * (t + 1))
                if (area[0] + area[2]) < x_in.shape[2]:
                    for t in range(rr):
                        mult[:,:,area[0] - 1 - t:area[0] - t,:] *= ((1.0/rr) * (t + 1))
                if area[3] != 0:
                    for t in range(rr):
                        mult[:,:,:,t:1+t] *= ((1.0/rr) * (t + 1))
                if (area[1] + area[3]) < x_in.shape[3]:
                    for t in range(rr):
                        mult[:,:,:,area[1] - 1 - t:area[1] - t] *= ((1.0/rr) * (t + 1))

comfyanonymous's avatar
comfyanonymous committed
69
70
            conditionning = {}
            conditionning['c_crossattn'] = cond[0]
71
72
73
74
75
76
77
78
79

            if 'concat' in cond[1]:
                cond_concat_in = cond[1]['concat']
                if cond_concat_in is not None and len(cond_concat_in) > 0:
                    cropped = []
                    for x in cond_concat_in:
                        cr = x[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
                        cropped.append(cr)
                    conditionning['c_concat'] = torch.cat(cropped, dim=1)
comfyanonymous's avatar
comfyanonymous committed
80

81
82
83
            if adm_cond is not None:
                conditionning['c_adm'] = adm_cond

comfyanonymous's avatar
comfyanonymous committed
84
85
86
            control = None
            if 'control' in cond[1]:
                control = cond[1]['control']
87
88
89
90
91
92
93
94

            patches = None
            if 'gligen' in cond[1]:
                gligen = cond[1]['gligen']
                patches = {}
                gligen_type = gligen[0]
                gligen_model = gligen[1]
                if gligen_type == "position":
comfyanonymous's avatar
comfyanonymous committed
95
                    gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
96
                else:
comfyanonymous's avatar
comfyanonymous committed
97
                    gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)
98
99
100
101

                patches['middle_patch'] = [gligen_patch]

            return (input_x, mult, conditionning, area, control, patches)
comfyanonymous's avatar
comfyanonymous committed
102
103

        def cond_equal_size(c1, c2):
comfyanonymous's avatar
comfyanonymous committed
104
105
            if c1 is c2:
                return True
comfyanonymous's avatar
comfyanonymous committed
106
107
108
            if c1.keys() != c2.keys():
                return False
            if 'c_crossattn' in c1:
comfyanonymous's avatar
comfyanonymous committed
109
110
111
112
113
114
115
116
117
118
                s1 = c1['c_crossattn'].shape
                s2 = c2['c_crossattn'].shape
                if s1 != s2:
                    if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen
                        return False

                    mult_min = lcm(s1[1], s2[1])
                    diff = mult_min // min(s1[1], s2[1])
                    if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
                        return False
comfyanonymous's avatar
comfyanonymous committed
119
120
121
            if 'c_concat' in c1:
                if c1['c_concat'].shape != c2['c_concat'].shape:
                    return False
122
123
124
            if 'c_adm' in c1:
                if c1['c_adm'].shape != c2['c_adm'].shape:
                    return False
comfyanonymous's avatar
comfyanonymous committed
125
126
            return True

comfyanonymous's avatar
comfyanonymous committed
127
128
129
        def can_concat_cond(c1, c2):
            if c1[0].shape != c2[0].shape:
                return False
130
131

            #control
comfyanonymous's avatar
comfyanonymous committed
132
133
134
135
136
137
            if (c1[4] is None) != (c2[4] is None):
                return False
            if c1[4] is not None:
                if c1[4] is not c2[4]:
                    return False

138
139
140
141
142
143
144
            #patches
            if (c1[5] is None) != (c2[5] is None):
                return False
            if (c1[5] is not None):
                if c1[5] is not c2[5]:
                    return False

comfyanonymous's avatar
comfyanonymous committed
145
146
            return cond_equal_size(c1[2], c2[2])

comfyanonymous's avatar
comfyanonymous committed
147
148
149
        def cond_cat(c_list):
            c_crossattn = []
            c_concat = []
150
            c_adm = []
comfyanonymous's avatar
comfyanonymous committed
151
            crossattn_max_len = 0
comfyanonymous's avatar
comfyanonymous committed
152
153
            for x in c_list:
                if 'c_crossattn' in x:
comfyanonymous's avatar
comfyanonymous committed
154
155
156
157
158
159
                    c = x['c_crossattn']
                    if crossattn_max_len == 0:
                        crossattn_max_len = c.shape[1]
                    else:
                        crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
                    c_crossattn.append(c)
comfyanonymous's avatar
comfyanonymous committed
160
161
                if 'c_concat' in x:
                    c_concat.append(x['c_concat'])
162
163
                if 'c_adm' in x:
                    c_adm.append(x['c_adm'])
comfyanonymous's avatar
comfyanonymous committed
164
            out = {}
comfyanonymous's avatar
comfyanonymous committed
165
166
167
168
169
170
171
            c_crossattn_out = []
            for c in c_crossattn:
                if c.shape[1] < crossattn_max_len:
                    c = c.repeat(1, crossattn_max_len // c.shape[1], 1) #padding with repeat doesn't change result
                c_crossattn_out.append(c)

            if len(c_crossattn_out) > 0:
172
                out['c_crossattn'] = torch.cat(c_crossattn_out)
comfyanonymous's avatar
comfyanonymous committed
173
            if len(c_concat) > 0:
174
                out['c_concat'] = torch.cat(c_concat)
175
176
            if len(c_adm) > 0:
                out['c_adm'] = torch.cat(c_adm)
comfyanonymous's avatar
comfyanonymous committed
177
178
            return out

179
        def calc_cond_uncond_batch(model_function, cond, uncond, x_in, timestep, max_total_area, model_options):
comfyanonymous's avatar
comfyanonymous committed
180
181
            out_cond = torch.zeros_like(x_in)
            out_count = torch.ones_like(x_in)/100000.0
182
183
184
185
186
187

            out_uncond = torch.zeros_like(x_in)
            out_uncond_count = torch.ones_like(x_in)/100000.0

            COND = 0
            UNCOND = 1
comfyanonymous's avatar
comfyanonymous committed
188

189
            to_run = []
comfyanonymous's avatar
comfyanonymous committed
190
            for x in cond:
191
                p = get_area_and_mult(x, x_in, timestep)
192
                if p is None:
comfyanonymous's avatar
comfyanonymous committed
193
                    continue
194
195

                to_run += [(p, COND)]
196
197
            if uncond is not None:
                for x in uncond:
198
                    p = get_area_and_mult(x, x_in, timestep)
199
200
                    if p is None:
                        continue
201

202
                    to_run += [(p, UNCOND)]
203
204
205
206

            while len(to_run) > 0:
                first = to_run[0]
                first_shape = first[0][0].shape
207
                to_batch_temp = []
208
                for x in range(len(to_run)):
comfyanonymous's avatar
comfyanonymous committed
209
210
                    if can_concat_cond(to_run[x][0], first[0]):
                        to_batch_temp += [x]
211
212
213
214
215
216
217
218
219

                to_batch_temp.reverse()
                to_batch = to_batch_temp[:1]

                for i in range(1, len(to_batch_temp) + 1):
                    batch_amount = to_batch_temp[:len(to_batch_temp)//i]
                    if (len(batch_amount) * first_shape[0] * first_shape[2] * first_shape[3] < max_total_area):
                        to_batch = batch_amount
                        break
220
221
222
223
224
225

                input_x = []
                mult = []
                c = []
                cond_or_uncond = []
                area = []
comfyanonymous's avatar
comfyanonymous committed
226
                control = None
227
                patches = None
228
229
230
231
232
233
234
235
                for x in to_batch:
                    o = to_run.pop(x)
                    p = o[0]
                    input_x += [p[0]]
                    mult += [p[1]]
                    c += [p[2]]
                    area += [p[3]]
                    cond_or_uncond += [o[1]]
comfyanonymous's avatar
comfyanonymous committed
236
                    control = p[4]
237
                    patches = p[5]
238
239
240

                batch_chunks = len(cond_or_uncond)
                input_x = torch.cat(input_x)
comfyanonymous's avatar
comfyanonymous committed
241
                c = cond_cat(c)
comfyanonymous's avatar
comfyanonymous committed
242
                timestep_ = torch.cat([timestep] * batch_chunks)
243

comfyanonymous's avatar
comfyanonymous committed
244
                if control is not None:
245
                    c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
comfyanonymous's avatar
comfyanonymous committed
246

247
                transformer_options = {}
248
                if 'transformer_options' in model_options:
249
250
251
                    transformer_options = model_options['transformer_options'].copy()

                if patches is not None:
252
253
254
255
256
257
258
259
260
                    if "patches" in transformer_options:
                        cur_patches = transformer_options["patches"].copy()
                        for p in patches:
                            if p in cur_patches:
                                cur_patches[p] = cur_patches[p] + patches[p]
                            else:
                                cur_patches[p] = patches[p]
                    else:
                        transformer_options["patches"] = patches
261

262
                transformer_options["cond_or_uncond"] = cond_or_uncond[:]
263
                c['transformer_options'] = transformer_options
264

265
266
267
268
                if 'model_function_wrapper' in model_options:
                    output = model_options['model_function_wrapper'](model_function, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
                else:
                    output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
comfyanonymous's avatar
comfyanonymous committed
269
                del input_x
270
271
272
273
274
275
276
277

                for o in range(batch_chunks):
                    if cond_or_uncond[o] == COND:
                        out_cond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
                        out_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
                    else:
                        out_uncond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
                        out_uncond_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
comfyanonymous's avatar
comfyanonymous committed
278
279
280
281
                del mult

            out_cond /= out_count
            del out_count
282
283
284
285
            out_uncond /= out_uncond_count
            del out_uncond_count

            return out_cond, out_uncond
comfyanonymous's avatar
comfyanonymous committed
286
287


288
        max_total_area = model_management.maximum_batch_area()
289
290
291
        if math.isclose(cond_scale, 1.0):
            uncond = None

292
        cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, model_options)
293
        if "sampler_cfg_function" in model_options:
294
295
            args = {"cond": cond, "uncond": uncond, "cond_scale": cond_scale, "timestep": timestep}
            return model_options["sampler_cfg_function"](args)
296
297
        else:
            return uncond + (cond - uncond) * cond_scale
comfyanonymous's avatar
comfyanonymous committed
298

comfyanonymous's avatar
comfyanonymous committed
299
300
301
302
303
304
305
306
307
308
309
310
311
312

class CompVisVDenoiser(k_diffusion_external.DiscreteVDDPMDenoiser):
    def __init__(self, model, quantize=False, device='cpu'):
        super().__init__(model, model.alphas_cumprod, quantize=quantize)

    def get_v(self, x, t, cond, **kwargs):
        return self.inner_model.apply_model(x, t, cond, **kwargs)


class CFGNoisePredictor(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
        self.alphas_cumprod = model.alphas_cumprod
313
314
    def apply_model(self, x, timestep, cond, uncond, cond_scale, model_options={}, seed=None):
        out = sampling_function(self.inner_model.apply_model, x, timestep, uncond, cond, cond_scale, model_options=model_options, seed=seed)
comfyanonymous's avatar
comfyanonymous committed
315
316
317
318
        return out


class KSamplerX0Inpaint(torch.nn.Module):
319
320
321
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
322
    def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
323
324
        if denoise_mask is not None:
            latent_mask = 1. - denoise_mask
325
            x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
326
        out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
327
328
329
330
331
332
        if denoise_mask is not None:
            out *= denoise_mask

        if denoise_mask is not None:
            out += self.latent_image * latent_mask
        return out
333

comfyanonymous's avatar
comfyanonymous committed
334
335
336
337
338
339
340
341
def simple_scheduler(model, steps):
    sigs = []
    ss = len(model.sigmas) / steps
    for x in range(steps):
        sigs += [float(model.sigmas[-(1 + int(x * ss))])]
    sigs += [0.0]
    return torch.FloatTensor(sigs)

comfyanonymous's avatar
comfyanonymous committed
342
343
344
345
def ddim_scheduler(model, steps):
    sigs = []
    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)
    for x in range(len(ddim_timesteps) - 1, -1, -1):
346
347
348
349
        ts = ddim_timesteps[x]
        if ts > 999:
            ts = 999
        sigs.append(model.t_to_sigma(torch.tensor(ts)))
comfyanonymous's avatar
comfyanonymous committed
350
351
352
    sigs += [0.0]
    return torch.FloatTensor(sigs)

353
354
355
356
357
358
359
360
361
362
363
def sgm_scheduler(model, steps):
    sigs = []
    timesteps = torch.linspace(model.inner_model.inner_model.num_timesteps - 1, 0, steps + 1)[:-1].type(torch.int)
    for x in range(len(timesteps)):
        ts = timesteps[x]
        if ts > 999:
            ts = 999
        sigs.append(model.t_to_sigma(torch.tensor(ts)))
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def get_mask_aabb(masks):
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device, dtype=torch.int)

    b = masks.shape[0]

    bounding_boxes = torch.zeros((b, 4), device=masks.device, dtype=torch.int)
    is_empty = torch.zeros((b), device=masks.device, dtype=torch.bool)
    for i in range(b):
        mask = masks[i]
        if mask.numel() == 0:
            continue
        if torch.max(mask != 0) == False:
            is_empty[i] = True
            continue
        y, x = torch.where(mask)
        bounding_boxes[i, 0] = torch.min(x)
        bounding_boxes[i, 1] = torch.min(y)
        bounding_boxes[i, 2] = torch.max(x)
        bounding_boxes[i, 3] = torch.max(y)

    return bounding_boxes, is_empty

387
def resolve_areas_and_cond_masks(conditions, h, w, device):
Jacob Segal's avatar
Jacob Segal committed
388
389
390
391
    # We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes.
    # While we're doing this, we can also resolve the mask device and scaling for performance reasons
    for i in range(len(conditions)):
        c = conditions[i]
392
393
394
395
396
397
398
399
400
        if 'area' in c[1]:
            area = c[1]['area']
            if area[0] == "percentage":
                modified = c[1].copy()
                area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
                modified['area'] = area
                c = [c[0], modified]
                conditions[i] = c

Jacob Segal's avatar
Jacob Segal committed
401
402
403
404
405
406
        if 'mask' in c[1]:
            mask = c[1]['mask']
            mask = mask.to(device=device)
            modified = c[1].copy()
            if len(mask.shape) == 2:
                mask = mask.unsqueeze(0)
mara's avatar
mara committed
407
            if mask.shape[1] != h or mask.shape[2] != w:
Jacob Segal's avatar
Jacob Segal committed
408
409
                mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1)

Jacob Segal's avatar
Jacob Segal committed
410
            if modified.get("set_area_to_bounds", False):
Jacob Segal's avatar
Jacob Segal committed
411
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
Jacob Segal's avatar
Jacob Segal committed
412
413
414
415
                boxes, is_empty = get_mask_aabb(bounds)
                if is_empty[0]:
                    # Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway)
                    modified['area'] = (8, 8, 0, 0)
Jacob Segal's avatar
Jacob Segal committed
416
                else:
Jacob Segal's avatar
Jacob Segal committed
417
                    box = boxes[0]
Jacob Segal's avatar
Jacob Segal committed
418
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
419
420
                    H = max(8, H)
                    W = max(8, W)
Jacob Segal's avatar
Jacob Segal committed
421
422
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area
Jacob Segal's avatar
Jacob Segal committed
423
424
425
426

            modified['mask'] = mask
            conditions[i] = [c[0], modified]

comfyanonymous's avatar
comfyanonymous committed
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
def create_cond_with_same_area_if_none(conds, c):
    if 'area' not in c[1]:
        return

    c_area = c[1]['area']
    smallest = None
    for x in conds:
        if 'area' in x[1]:
            a = x[1]['area']
            if c_area[2] >= a[2] and c_area[3] >= a[3]:
                if a[0] + a[2] >= c_area[0] + c_area[2]:
                    if a[1] + a[3] >= c_area[1] + c_area[3]:
                        if smallest is None:
                            smallest = x
                        elif 'area' not in smallest[1]:
                            smallest = x
                        else:
                            if smallest[1]['area'][0] * smallest[1]['area'][1] > a[0] * a[1]:
                                smallest = x
        else:
            if smallest is None:
                smallest = x
    if smallest is None:
        return
    if 'area' in smallest[1]:
        if smallest[1]['area'] == c_area:
            return
    n = c[1].copy()
    conds += [[smallest[0], n]]
comfyanonymous's avatar
comfyanonymous committed
456

457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
def calculate_start_end_timesteps(model, conds):
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
        if 'start_percent' in x[1]:
            timestep_start = model.sigma_to_t(model.t_to_sigma(torch.tensor(x[1]['start_percent'] * 999.0)))
        if 'end_percent' in x[1]:
            timestep_end = model.sigma_to_t(model.t_to_sigma(torch.tensor(x[1]['end_percent'] * 999.0)))

        if (timestep_start is not None) or (timestep_end is not None):
            n = x[1].copy()
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
            conds[t] = [x[0], n]

476
477
478
479
480
481
482
483
def pre_run_control(model, conds):
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
        percent_to_timestep_function = lambda a: model.sigma_to_t(model.t_to_sigma(torch.tensor(a) * 999.0))
        if 'control' in x[1]:
comfyanonymous's avatar
comfyanonymous committed
484
            x[1]['control'].pre_run(model.inner_model.inner_model, percent_to_timestep_function)
485

486
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
comfyanonymous's avatar
comfyanonymous committed
487
488
489
490
491
492
493
    cond_cnets = []
    cond_other = []
    uncond_cnets = []
    uncond_other = []
    for t in range(len(conds)):
        x = conds[t]
        if 'area' not in x[1]:
494
495
            if name in x[1] and x[1][name] is not None:
                cond_cnets.append(x[1][name])
comfyanonymous's avatar
comfyanonymous committed
496
497
498
499
500
            else:
                cond_other.append((x, t))
    for t in range(len(uncond)):
        x = uncond[t]
        if 'area' not in x[1]:
501
502
            if name in x[1] and x[1][name] is not None:
                uncond_cnets.append(x[1][name])
comfyanonymous's avatar
comfyanonymous committed
503
504
505
506
507
508
509
510
511
            else:
                uncond_other.append((x, t))

    if len(uncond_cnets) > 0:
        return

    for x in range(len(cond_cnets)):
        temp = uncond_other[x % len(uncond_other)]
        o = temp[0]
512
        if name in o[1] and o[1][name] is not None:
comfyanonymous's avatar
comfyanonymous committed
513
            n = o[1].copy()
514
            n[name] = uncond_fill_func(cond_cnets, x)
comfyanonymous's avatar
comfyanonymous committed
515
516
517
            uncond += [[o[0], n]]
        else:
            n = o[1].copy()
518
            n[name] = uncond_fill_func(cond_cnets, x)
comfyanonymous's avatar
comfyanonymous committed
519
520
            uncond[temp[1]] = [o[0], n]

521
def encode_adm(model, conds, batch_size, width, height, device, prompt_type):
522
523
    for t in range(len(conds)):
        x = conds[t]
comfyanonymous's avatar
comfyanonymous committed
524
        adm_out = None
525
526
        if 'adm' in x[1]:
            adm_out = x[1]["adm"]
527
        else:
528
            params = x[1].copy()
529
530
531
            params["width"] = params.get("width", width * 8)
            params["height"] = params.get("height", height * 8)
            params["prompt_type"] = params.get("prompt_type", prompt_type)
532
            adm_out = model.encode_adm(device=device, **params)
533

comfyanonymous's avatar
comfyanonymous committed
534
535
        if adm_out is not None:
            x[1] = x[1].copy()
536
            x[1]["adm_encoded"] = comfy.utils.repeat_to_batch_size(adm_out, batch_size).to(device)
537
538
539

    return conds

540
def encode_cond(model_function, key, conds, device, **kwargs):
541
542
543
    for t in range(len(conds)):
        x = conds[t]
        params = x[1].copy()
544
        params["device"] = device
545
546
547
548
549
550
551
552
553
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
        if out is not None:
            x[1] = x[1].copy()
            x[1][key] = out
    return conds
554

comfyanonymous's avatar
comfyanonymous committed
555
556
557
558
559
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
560
        return math.isclose(float(model_wrap.sigma_max), float(sigmas[0]), rel_tol=1e-05)
comfyanonymous's avatar
comfyanonymous committed
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608

class DDIM(Sampler):
    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
        timesteps = []
        for s in range(sigmas.shape[0]):
            timesteps.insert(0, model_wrap.sigma_to_discrete_timestep(sigmas[s]))
        noise_mask = None
        if denoise_mask is not None:
            noise_mask = 1.0 - denoise_mask

        ddim_callback = None
        if callback is not None:
            total_steps = len(timesteps) - 1
            ddim_callback = lambda pred_x0, i: callback(i, pred_x0, None, total_steps)

        max_denoise = self.max_denoise(model_wrap, sigmas)

        ddim_sampler = DDIMSampler(model_wrap.inner_model.inner_model, device=noise.device)
        ddim_sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
        z_enc = ddim_sampler.stochastic_encode(latent_image, torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(noise.device), noise=noise, max_denoise=max_denoise)
        samples, _ = ddim_sampler.sample_custom(ddim_timesteps=timesteps,
                                                batch_size=noise.shape[0],
                                                shape=noise.shape[1:],
                                                verbose=False,
                                                eta=0.0,
                                                x_T=z_enc,
                                                x0=latent_image,
                                                img_callback=ddim_callback,
                                                denoise_function=model_wrap.predict_eps_discrete_timestep,
                                                extra_args=extra_args,
                                                mask=noise_mask,
                                                to_zero=sigmas[-1]==0,
                                                end_step=sigmas.shape[0] - 1,
                                                disable_pbar=disable_pbar)
        return samples

class UNIPC(Sampler):
    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
        return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)

class UNIPCBH2(Sampler):
    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
        return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)

KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm"]

comfyanonymous's avatar
comfyanonymous committed
609
def ksampler(sampler_name, extra_options={}):
comfyanonymous's avatar
comfyanonymous committed
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
    class KSAMPLER(Sampler):
        def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
            extra_args["denoise_mask"] = denoise_mask
            model_k = KSamplerX0Inpaint(model_wrap)
            model_k.latent_image = latent_image
            model_k.noise = noise

            if self.max_denoise(model_wrap, sigmas):
                noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
            else:
                noise = noise * sigmas[0]

            k_callback = None
            total_steps = len(sigmas) - 1
            if callback is not None:
                k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)

            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]

            if latent_image is not None:
                noise += latent_image
            if sampler_name == "dpm_fast":
                samples = k_diffusion_sampling.sample_dpm_fast(model_k, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
            elif sampler_name == "dpm_adaptive":
                samples = k_diffusion_sampling.sample_dpm_adaptive(model_k, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=k_callback, disable=disable_pbar)
            else:
comfyanonymous's avatar
comfyanonymous committed
638
                samples = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **extra_options)
comfyanonymous's avatar
comfyanonymous committed
639
640
641
            return samples
    return KSAMPLER

comfyanonymous's avatar
comfyanonymous committed
642
643
644
645
646
647
648
def wrap_model(model):
    model_denoise = CFGNoisePredictor(model)
    if model.model_type == model_base.ModelType.V_PREDICTION:
        model_wrap = CompVisVDenoiser(model_denoise, quantize=True)
    else:
        model_wrap = k_diffusion_external.CompVisDenoiser(model_denoise, quantize=True)
    return model_wrap
comfyanonymous's avatar
comfyanonymous committed
649
650
651
652
653
654
655
656

def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
    positive = positive[:]
    negative = negative[:]

    resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
    resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)

comfyanonymous's avatar
comfyanonymous committed
657
    model_wrap = wrap_model(model)
comfyanonymous's avatar
comfyanonymous committed
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672

    calculate_start_end_timesteps(model_wrap, negative)
    calculate_start_end_timesteps(model_wrap, positive)

    #make sure each cond area has an opposite one with the same area
    for c in positive:
        create_cond_with_same_area_if_none(negative, c)
    for c in negative:
        create_cond_with_same_area_if_none(positive, c)

    pre_run_control(model_wrap, negative + positive)

    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])
    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])

673
674
675
    if latent_image is not None:
        latent_image = model.process_latent_in(latent_image)

comfyanonymous's avatar
comfyanonymous committed
676
677
678
679
    if model.is_adm():
        positive = encode_adm(model, positive, noise.shape[0], noise.shape[3], noise.shape[2], device, "positive")
        negative = encode_adm(model, negative, noise.shape[0], noise.shape[3], noise.shape[2], device, "negative")

680
    if hasattr(model, 'cond_concat'):
681
682
        positive = encode_cond(model.cond_concat, "concat", positive, device, noise=noise, latent_image=latent_image, denoise_mask=denoise_mask)
        negative = encode_cond(model.cond_concat, "concat", negative, device, noise=noise, latent_image=latent_image, denoise_mask=denoise_mask)
comfyanonymous's avatar
comfyanonymous committed
683
684
685
686
687
688

    extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}

    samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
    return model.process_latent_out(samples.to(torch.float32))

comfyanonymous's avatar
comfyanonymous committed
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]

def calculate_sigmas_scheduler(model, scheduler_name, steps):
    model_wrap = wrap_model(model)
    if scheduler_name == "karras":
        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_wrap.sigma_min), sigma_max=float(model_wrap.sigma_max))
    elif scheduler_name == "exponential":
        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_wrap.sigma_min), sigma_max=float(model_wrap.sigma_max))
    elif scheduler_name == "normal":
        sigmas = model_wrap.get_sigmas(steps)
    elif scheduler_name == "simple":
        sigmas = simple_scheduler(model_wrap, steps)
    elif scheduler_name == "ddim_uniform":
        sigmas = ddim_scheduler(model_wrap, steps)
    elif scheduler_name == "sgm_uniform":
        sigmas = sgm_scheduler(model_wrap, steps)
    else:
        print("error invalid scheduler", self.scheduler)
    return sigmas

710
711
712
713
714
715
716
717
718
719
720
def sampler_class(name):
    if name == "uni_pc":
        sampler = UNIPC
    elif name == "uni_pc_bh2":
        sampler = UNIPCBH2
    elif name == "ddim":
        sampler = DDIM
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
721
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
722
723
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
comfyanonymous's avatar
comfyanonymous committed
724

725
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
726
727
728
729
730
731
732
733
734
        self.model = model
        self.device = device
        if scheduler not in self.SCHEDULERS:
            scheduler = self.SCHEDULERS[0]
        if sampler not in self.SAMPLERS:
            sampler = self.SAMPLERS[0]
        self.scheduler = scheduler
        self.sampler = sampler
        self.set_steps(steps, denoise)
735
        self.denoise = denoise
736
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
737

comfyanonymous's avatar
comfyanonymous committed
738
739
740
741
742
743
744
745
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
        if self.sampler in ['dpm_2', 'dpm_2_ancestral']:
            steps += 1
            discard_penultimate_sigma = True

comfyanonymous's avatar
comfyanonymous committed
746
        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
747
748
749
750
751

        if discard_penultimate_sigma:
            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
        return sigmas

comfyanonymous's avatar
comfyanonymous committed
752
753
    def set_steps(self, steps, denoise=None):
        self.steps = steps
754
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
755
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
756
757
        else:
            new_steps = int(steps/denoise)
comfyanonymous's avatar
comfyanonymous committed
758
            sigmas = self.calculate_sigmas(new_steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
759
760
            self.sigmas = sigmas[-(steps + 1):]

761
    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):
762
763
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
764

comfyanonymous's avatar
comfyanonymous committed
765
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
766
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
767
768
769
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
770
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
771
772
773
774
775
776
777
            if start_step < (len(sigmas) - 1):
                sigmas = sigmas[start_step:]
            else:
                if latent_image is not None:
                    return latent_image
                else:
                    return torch.zeros_like(noise)
comfyanonymous's avatar
comfyanonymous committed
778

779
        sampler = sampler_class(self.sampler)
780

comfyanonymous's avatar
comfyanonymous committed
781
        return sample(self.model, noise, positive, negative, cfg, self.device, sampler(), sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)