samplers.py 27.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
import enum
6
from comfy import model_management
comfyanonymous's avatar
comfyanonymous committed
7
import math
8
from comfy import model_base
9
import comfy.utils
10
import comfy.conds
11
12


comfyanonymous's avatar
comfyanonymous committed
13
#The main sampling function shared by all the samplers
comfyanonymous's avatar
comfyanonymous committed
14
#Returns denoised
15
def sampling_function(model_function, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
16
        def get_area_and_mult(conds, x_in, timestep_in):
17
18
            area = (x_in.shape[2], x_in.shape[3], 0, 0)
            strength = 1.0
19
20
21

            if 'timestep_start' in conds:
                timestep_start = conds['timestep_start']
22
                if timestep_in[0] > timestep_start:
23
                    return None
24
25
            if 'timestep_end' in conds:
                timestep_end = conds['timestep_end']
26
                if timestep_in[0] < timestep_end:
27
                    return None
28
29
30
31
            if 'area' in conds:
                area = conds['area']
            if 'strength' in conds:
                strength = conds['strength']
32

33
            input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
34
            if 'mask' in conds:
Jacob Segal's avatar
Jacob Segal committed
35
36
                # Scale the mask to the size of the input
                # The mask should have been resized as we began the sampling process
37
                mask_strength = 1.0
38
39
40
                if "mask_strength" in conds:
                    mask_strength = conds["mask_strength"]
                mask = conds['mask']
Jacob Segal's avatar
Jacob Segal committed
41
42
                assert(mask.shape[1] == x_in.shape[2])
                assert(mask.shape[2] == x_in.shape[3])
43
                mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
Jacob Segal's avatar
Jacob Segal committed
44
                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
45
46
47
48
            else:
                mask = torch.ones_like(input_x)
            mult = mask * strength

49
            if 'mask' not in conds:
Jacob Segal's avatar
Jacob Segal committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
                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
64
            conditionning = {}
65
66
67
            model_conds = conds["model_conds"]
            for c in model_conds:
                conditionning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)
68

comfyanonymous's avatar
comfyanonymous committed
69
            control = None
70
71
            if 'control' in conds:
                control = conds['control']
72
73

            patches = None
74
75
            if 'gligen' in conds:
                gligen = conds['gligen']
76
77
78
79
                patches = {}
                gligen_type = gligen[0]
                gligen_model = gligen[1]
                if gligen_type == "position":
comfyanonymous's avatar
comfyanonymous committed
80
                    gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
81
                else:
comfyanonymous's avatar
comfyanonymous committed
82
                    gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)
83
84
85
86

                patches['middle_patch'] = [gligen_patch]

            return (input_x, mult, conditionning, area, control, patches)
comfyanonymous's avatar
comfyanonymous committed
87
88

        def cond_equal_size(c1, c2):
comfyanonymous's avatar
comfyanonymous committed
89
90
            if c1 is c2:
                return True
comfyanonymous's avatar
comfyanonymous committed
91
92
            if c1.keys() != c2.keys():
                return False
93
94
            for k in c1:
                if not c1[k].can_concat(c2[k]):
95
                    return False
comfyanonymous's avatar
comfyanonymous committed
96
97
            return True

comfyanonymous's avatar
comfyanonymous committed
98
99
100
        def can_concat_cond(c1, c2):
            if c1[0].shape != c2[0].shape:
                return False
101
102

            #control
comfyanonymous's avatar
comfyanonymous committed
103
104
105
106
107
108
            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

109
110
111
112
113
114
115
            #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
116
117
            return cond_equal_size(c1[2], c2[2])

comfyanonymous's avatar
comfyanonymous committed
118
119
120
        def cond_cat(c_list):
            c_crossattn = []
            c_concat = []
121
            c_adm = []
comfyanonymous's avatar
comfyanonymous committed
122
            crossattn_max_len = 0
123
124

            temp = {}
comfyanonymous's avatar
comfyanonymous committed
125
            for x in c_list:
126
127
128
129
130
                for k in x:
                    cur = temp.get(k, [])
                    cur.append(x[k])
                    temp[k] = cur

comfyanonymous's avatar
comfyanonymous committed
131
            out = {}
132
133
134
135
            for k in temp:
                conds = temp[k]
                out[k] = conds[0].concat(conds[1:])

comfyanonymous's avatar
comfyanonymous committed
136
137
            return out

138
        def calc_cond_uncond_batch(model_function, cond, uncond, x_in, timestep, max_total_area, model_options):
comfyanonymous's avatar
comfyanonymous committed
139
140
            out_cond = torch.zeros_like(x_in)
            out_count = torch.ones_like(x_in)/100000.0
141
142
143
144
145
146

            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
147

148
            to_run = []
comfyanonymous's avatar
comfyanonymous committed
149
            for x in cond:
150
                p = get_area_and_mult(x, x_in, timestep)
151
                if p is None:
comfyanonymous's avatar
comfyanonymous committed
152
                    continue
153
154

                to_run += [(p, COND)]
155
156
            if uncond is not None:
                for x in uncond:
157
                    p = get_area_and_mult(x, x_in, timestep)
158
159
                    if p is None:
                        continue
160

161
                    to_run += [(p, UNCOND)]
162
163
164
165

            while len(to_run) > 0:
                first = to_run[0]
                first_shape = first[0][0].shape
166
                to_batch_temp = []
167
                for x in range(len(to_run)):
comfyanonymous's avatar
comfyanonymous committed
168
169
                    if can_concat_cond(to_run[x][0], first[0]):
                        to_batch_temp += [x]
170
171
172
173
174
175
176
177
178

                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
179
180
181
182
183
184

                input_x = []
                mult = []
                c = []
                cond_or_uncond = []
                area = []
comfyanonymous's avatar
comfyanonymous committed
185
                control = None
186
                patches = None
187
188
189
190
191
192
193
194
                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
195
                    control = p[4]
196
                    patches = p[5]
197
198
199

                batch_chunks = len(cond_or_uncond)
                input_x = torch.cat(input_x)
comfyanonymous's avatar
comfyanonymous committed
200
                c = cond_cat(c)
comfyanonymous's avatar
comfyanonymous committed
201
                timestep_ = torch.cat([timestep] * batch_chunks)
202

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

206
                transformer_options = {}
207
                if 'transformer_options' in model_options:
208
209
210
                    transformer_options = model_options['transformer_options'].copy()

                if patches is not None:
211
212
213
214
215
216
217
218
219
                    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
220

221
                transformer_options["cond_or_uncond"] = cond_or_uncond[:]
222
                c['transformer_options'] = transformer_options
223

224
225
226
227
                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
228
                del input_x
229
230
231
232
233
234
235
236

                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
237
238
239
240
                del mult

            out_cond /= out_count
            del out_count
241
242
243
244
            out_uncond /= out_uncond_count
            del out_uncond_count

            return out_cond, out_uncond
comfyanonymous's avatar
comfyanonymous committed
245
246


247
        max_total_area = model_management.maximum_batch_area()
248
249
250
        if math.isclose(cond_scale, 1.0):
            uncond = None

251
        cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, model_options)
252
        if "sampler_cfg_function" in model_options:
253
254
            args = {"cond": cond, "uncond": uncond, "cond_scale": cond_scale, "timestep": timestep}
            return model_options["sampler_cfg_function"](args)
255
256
        else:
            return uncond + (cond - uncond) * cond_scale
comfyanonymous's avatar
comfyanonymous committed
257

comfyanonymous's avatar
comfyanonymous committed
258
259
260
261
class CFGNoisePredictor(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
262
263
    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
264
        return out
comfyanonymous's avatar
comfyanonymous committed
265
266
    def forward(self, *args, **kwargs):
        return self.apply_model(*args, **kwargs)
comfyanonymous's avatar
comfyanonymous committed
267
268

class KSamplerX0Inpaint(torch.nn.Module):
269
270
271
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
272
    def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
273
274
        if denoise_mask is not None:
            latent_mask = 1. - denoise_mask
275
            x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
276
        out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
277
278
279
280
281
282
        if denoise_mask is not None:
            out *= denoise_mask

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

comfyanonymous's avatar
comfyanonymous committed
284
def simple_scheduler(model, steps):
comfyanonymous's avatar
comfyanonymous committed
285
    s = model.model_sampling
comfyanonymous's avatar
comfyanonymous committed
286
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
287
    ss = len(s.sigmas) / steps
comfyanonymous's avatar
comfyanonymous committed
288
    for x in range(steps):
comfyanonymous's avatar
comfyanonymous committed
289
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
comfyanonymous's avatar
comfyanonymous committed
290
291
292
    sigs += [0.0]
    return torch.FloatTensor(sigs)

comfyanonymous's avatar
comfyanonymous committed
293
def ddim_scheduler(model, steps):
comfyanonymous's avatar
comfyanonymous committed
294
    s = model.model_sampling
comfyanonymous's avatar
comfyanonymous committed
295
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
296
297
298
299
300
301
    ss = len(s.sigmas) // steps
    x = 1
    while x < len(s.sigmas):
        sigs += [float(s.sigmas[x])]
        x += ss
    sigs = sigs[::-1]
comfyanonymous's avatar
comfyanonymous committed
302
303
304
    sigs += [0.0]
    return torch.FloatTensor(sigs)

comfyanonymous's avatar
comfyanonymous committed
305
306
307
308
309
310
311
312
313
314
def normal_scheduler(model, steps, sgm=False, floor=False):
    s = model.model_sampling
    start = s.timestep(s.sigma_max)
    end = s.timestep(s.sigma_min)

    if sgm:
        timesteps = torch.linspace(start, end, steps + 1)[:-1]
    else:
        timesteps = torch.linspace(start, end, steps)

315
316
317
    sigs = []
    for x in range(len(timesteps)):
        ts = timesteps[x]
comfyanonymous's avatar
comfyanonymous committed
318
        sigs.append(s.sigma(ts))
319
320
321
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
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

345
def resolve_areas_and_cond_masks(conditions, h, w, device):
Jacob Segal's avatar
Jacob Segal committed
346
347
348
349
    # 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]
350
351
        if 'area' in c:
            area = c['area']
352
            if area[0] == "percentage":
353
                modified = c.copy()
354
355
                area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
                modified['area'] = area
356
                c = modified
357
358
                conditions[i] = c

359
360
        if 'mask' in c:
            mask = c['mask']
Jacob Segal's avatar
Jacob Segal committed
361
            mask = mask.to(device=device)
362
            modified = c.copy()
Jacob Segal's avatar
Jacob Segal committed
363
364
            if len(mask.shape) == 2:
                mask = mask.unsqueeze(0)
mara's avatar
mara committed
365
            if mask.shape[1] != h or mask.shape[2] != w:
Jacob Segal's avatar
Jacob Segal committed
366
367
                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
368
            if modified.get("set_area_to_bounds", False):
Jacob Segal's avatar
Jacob Segal committed
369
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
Jacob Segal's avatar
Jacob Segal committed
370
371
372
373
                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
374
                else:
Jacob Segal's avatar
Jacob Segal committed
375
                    box = boxes[0]
Jacob Segal's avatar
Jacob Segal committed
376
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
377
378
                    H = max(8, H)
                    W = max(8, W)
Jacob Segal's avatar
Jacob Segal committed
379
380
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area
Jacob Segal's avatar
Jacob Segal committed
381
382

            modified['mask'] = mask
383
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
384

comfyanonymous's avatar
comfyanonymous committed
385
def create_cond_with_same_area_if_none(conds, c):
386
    if 'area' not in c:
comfyanonymous's avatar
comfyanonymous committed
387
388
        return

389
    c_area = c['area']
comfyanonymous's avatar
comfyanonymous committed
390
391
    smallest = None
    for x in conds:
392
393
        if 'area' in x:
            a = x['area']
comfyanonymous's avatar
comfyanonymous committed
394
395
396
397
398
            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
399
                        elif 'area' not in smallest:
comfyanonymous's avatar
comfyanonymous committed
400
401
                            smallest = x
                        else:
402
                            if smallest['area'][0] * smallest['area'][1] > a[0] * a[1]:
comfyanonymous's avatar
comfyanonymous committed
403
404
405
406
407
408
                                smallest = x
        else:
            if smallest is None:
                smallest = x
    if smallest is None:
        return
409
410
    if 'area' in smallest:
        if smallest['area'] == c_area:
comfyanonymous's avatar
comfyanonymous committed
411
            return
412
413
414
415

    out = c.copy()
    out['model_conds'] = smallest['model_conds'].copy() #TODO: which fields should be copied?
    conds += [out]
comfyanonymous's avatar
comfyanonymous committed
416

417
def calculate_start_end_timesteps(model, conds):
418
    s = model.model_sampling
419
420
421
422
423
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
424
        if 'start_percent' in x:
425
            timestep_start = s.percent_to_sigma(x['start_percent'])
426
        if 'end_percent' in x:
427
            timestep_end = s.percent_to_sigma(x['end_percent'])
428
429

        if (timestep_start is not None) or (timestep_end is not None):
430
            n = x.copy()
431
432
433
434
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
435
            conds[t] = n
436

437
def pre_run_control(model, conds):
438
    s = model.model_sampling
439
440
441
442
443
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
444
        percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
445
        if 'control' in x:
446
            x['control'].pre_run(model, percent_to_timestep_function)
447

448
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
comfyanonymous's avatar
comfyanonymous committed
449
450
451
452
453
454
    cond_cnets = []
    cond_other = []
    uncond_cnets = []
    uncond_other = []
    for t in range(len(conds)):
        x = conds[t]
455
456
457
        if 'area' not in x:
            if name in x and x[name] is not None:
                cond_cnets.append(x[name])
comfyanonymous's avatar
comfyanonymous committed
458
459
460
461
            else:
                cond_other.append((x, t))
    for t in range(len(uncond)):
        x = uncond[t]
462
463
464
        if 'area' not in x:
            if name in x and x[name] is not None:
                uncond_cnets.append(x[name])
comfyanonymous's avatar
comfyanonymous committed
465
466
467
468
469
470
471
472
473
            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]
474
475
        if name in o and o[name] is not None:
            n = o.copy()
476
            n[name] = uncond_fill_func(cond_cnets, x)
477
            uncond += [n]
comfyanonymous's avatar
comfyanonymous committed
478
        else:
479
            n = o.copy()
480
            n[name] = uncond_fill_func(cond_cnets, x)
481
            uncond[temp[1]] = n
482

483
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
484
485
    for t in range(len(conds)):
        x = conds[t]
486
        params = x.copy()
487
        params["device"] = device
488
489
490
491
        params["noise"] = noise
        params["width"] = params.get("width", noise.shape[3] * 8)
        params["height"] = params.get("height", noise.shape[2] * 8)
        params["prompt_type"] = params.get("prompt_type", prompt_type)
492
493
494
495
496
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
497
498
499
500
501
502
        x = x.copy()
        model_conds = x['model_conds'].copy()
        for k in out:
            model_conds[k] = out[k]
        x['model_conds'] = model_conds
        conds[t] = x
503
    return conds
504

comfyanonymous's avatar
comfyanonymous committed
505
506
507
508
509
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
510
511
512
        max_sigma = float(model_wrap.inner_model.model_sampling.sigma_max)
        sigma = float(sigmas[0])
        return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
comfyanonymous's avatar
comfyanonymous committed
513
514
515
516
517
518
519
520
521
522
523
524
525

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"]

526
def ksampler(sampler_name, extra_options={}, inpaint_options={}):
comfyanonymous's avatar
comfyanonymous committed
527
528
529
530
531
    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
532
533
534
535
536
            if inpaint_options.get("random", False): #TODO: Should this be the default?
                generator = torch.manual_seed(extra_args.get("seed", 41) + 1)
                model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
            else:
                model_k.noise = noise
comfyanonymous's avatar
comfyanonymous committed
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558

            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
559
                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
560
561
562
            return samples
    return KSAMPLER

comfyanonymous's avatar
comfyanonymous committed
563
564
def wrap_model(model):
    model_denoise = CFGNoisePredictor(model)
comfyanonymous's avatar
comfyanonymous committed
565
    return model_denoise
comfyanonymous's avatar
comfyanonymous committed
566
567
568
569
570
571
572
573

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
574
    model_wrap = wrap_model(model)
comfyanonymous's avatar
comfyanonymous committed
575

576
577
    calculate_start_end_timesteps(model, negative)
    calculate_start_end_timesteps(model, positive)
comfyanonymous's avatar
comfyanonymous committed
578
579
580
581
582
583
584

    #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)

585
    pre_run_control(model, negative + positive)
comfyanonymous's avatar
comfyanonymous committed
586

587
    apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
comfyanonymous's avatar
comfyanonymous committed
588
589
    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])

590
591
592
    if latent_image is not None:
        latent_image = model.process_latent_in(latent_image)

593
594
595
    if hasattr(model, 'extra_conds'):
        positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
        negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
comfyanonymous's avatar
comfyanonymous committed
596
597
598
599
600
601

    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
602
603
604
605
606
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):
    if scheduler_name == "karras":
comfyanonymous's avatar
comfyanonymous committed
607
        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
608
    elif scheduler_name == "exponential":
comfyanonymous's avatar
comfyanonymous committed
609
        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
610
    elif scheduler_name == "normal":
comfyanonymous's avatar
comfyanonymous committed
611
        sigmas = normal_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
612
    elif scheduler_name == "simple":
comfyanonymous's avatar
comfyanonymous committed
613
        sigmas = simple_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
614
    elif scheduler_name == "ddim_uniform":
comfyanonymous's avatar
comfyanonymous committed
615
        sigmas = ddim_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
616
    elif scheduler_name == "sgm_uniform":
comfyanonymous's avatar
comfyanonymous committed
617
        sigmas = normal_scheduler(model, steps, sgm=True)
comfyanonymous's avatar
comfyanonymous committed
618
619
620
621
    else:
        print("error invalid scheduler", self.scheduler)
    return sigmas

622
623
624
625
626
627
def sampler_class(name):
    if name == "uni_pc":
        sampler = UNIPC
    elif name == "uni_pc_bh2":
        sampler = UNIPCBH2
    elif name == "ddim":
628
        sampler = ksampler("euler", inpaint_options={"random": True})
629
630
631
632
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
633
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
634
635
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
comfyanonymous's avatar
comfyanonymous committed
636

637
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
638
639
640
641
642
643
644
645
646
        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)
647
        self.denoise = denoise
648
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
649

comfyanonymous's avatar
comfyanonymous committed
650
651
652
653
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
654
        if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
comfyanonymous's avatar
comfyanonymous committed
655
656
657
            steps += 1
            discard_penultimate_sigma = True

comfyanonymous's avatar
comfyanonymous committed
658
        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
659
660
661
662
663

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

comfyanonymous's avatar
comfyanonymous committed
664
665
    def set_steps(self, steps, denoise=None):
        self.steps = steps
666
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
667
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
668
669
        else:
            new_steps = int(steps/denoise)
comfyanonymous's avatar
comfyanonymous committed
670
            sigmas = self.calculate_sigmas(new_steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
671
672
            self.sigmas = sigmas[-(steps + 1):]

673
    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):
674
675
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
676

comfyanonymous's avatar
comfyanonymous committed
677
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
678
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
679
680
681
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
682
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
683
684
685
686
687
688
689
            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
690

691
        sampler = sampler_class(self.sampler)
692

comfyanonymous's avatar
comfyanonymous committed
693
        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)