samplers.py 30.3 KB
Newer Older
1
from .k_diffusion import sampling as k_diffusion_sampling
2
from .extra_samplers import uni_pc
comfyanonymous's avatar
comfyanonymous committed
3
import torch
comfyanonymous's avatar
comfyanonymous committed
4
import collections
5
from comfy import model_management
comfyanonymous's avatar
comfyanonymous committed
6
import math
7
import logging
8
import comfy.sampler_helpers
9

10
def get_area_and_mult(conds, x_in, timestep_in):
11
12
    dims = tuple(x_in.shape[2:])
    area = None
13
14
15
16
17
18
19
20
21
22
23
    strength = 1.0

    if 'timestep_start' in conds:
        timestep_start = conds['timestep_start']
        if timestep_in[0] > timestep_start:
            return None
    if 'timestep_end' in conds:
        timestep_end = conds['timestep_end']
        if timestep_in[0] < timestep_end:
            return None
    if 'area' in conds:
24
        area = list(conds['area'])
25
26
27
    if 'strength' in conds:
        strength = conds['strength']

28
29
30
31
32
33
    input_x = x_in
    if area is not None:
        for i in range(len(dims)):
            area[i] = min(input_x.shape[i + 2] - area[len(dims) + i], area[i])
            input_x = input_x.narrow(i + 2, area[len(dims) + i], area[i])

34
35
36
37
38
39
40
    if 'mask' in conds:
        # Scale the mask to the size of the input
        # The mask should have been resized as we began the sampling process
        mask_strength = 1.0
        if "mask_strength" in conds:
            mask_strength = conds["mask_strength"]
        mask = conds['mask']
41
42
43
44
45
46
47
48
        assert(mask.shape[1:] == x_in.shape[2:])

        mask = mask[:input_x.shape[0]]
        if area is not None:
            for i in range(len(dims)):
                mask = mask.narrow(i + 1, area[len(dims) + i], area[i])

        mask = mask * mask_strength
49
50
51
52
53
        mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
    else:
        mask = torch.ones_like(input_x)
    mult = mask * strength

54
    if 'mask' not in conds and area is not None:
55
        rr = 8
56
57
58
59
60
61
62
63
64
        for i in range(len(dims)):
            if area[len(dims) + i] != 0:
                for t in range(rr):
                    m = mult.narrow(i + 2, t, 1)
                    m *= ((1.0/rr) * (t + 1))
            if (area[i] + area[len(dims) + i]) < x_in.shape[i + 2]:
                for t in range(rr):
                    m = mult.narrow(i + 2, area[i] - 1 - t, 1)
                    m *= ((1.0/rr) * (t + 1))
65
66
67
68
69
70

    conditioning = {}
    model_conds = conds["model_conds"]
    for c in model_conds:
        conditioning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)

comfyanonymous's avatar
comfyanonymous committed
71
    control = conds.get('control', None)
72
73
74
75
76
77
78
79
80
81
82
83
84
85

    patches = None
    if 'gligen' in conds:
        gligen = conds['gligen']
        patches = {}
        gligen_type = gligen[0]
        gligen_model = gligen[1]
        if gligen_type == "position":
            gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
        else:
            gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)

        patches['middle_patch'] = [gligen_patch]

comfyanonymous's avatar
comfyanonymous committed
86
87
    cond_obj = collections.namedtuple('cond_obj', ['input_x', 'mult', 'conditioning', 'area', 'control', 'patches'])
    return cond_obj(input_x, mult, conditioning, area, control, patches)
88
89
90
91
92
93
94
95
96
97
98
99

def cond_equal_size(c1, c2):
    if c1 is c2:
        return True
    if c1.keys() != c2.keys():
        return False
    for k in c1:
        if not c1[k].can_concat(c2[k]):
            return False
    return True

def can_concat_cond(c1, c2):
comfyanonymous's avatar
comfyanonymous committed
100
    if c1.input_x.shape != c2.input_x.shape:
101
102
        return False

comfyanonymous's avatar
comfyanonymous committed
103
104
    def objects_concatable(obj1, obj2):
        if (obj1 is None) != (obj2 is None):
105
            return False
comfyanonymous's avatar
comfyanonymous committed
106
107
108
109
        if obj1 is not None:
            if obj1 is not obj2:
                return False
        return True
110

comfyanonymous's avatar
comfyanonymous committed
111
112
113
114
    if not objects_concatable(c1.control, c2.control):
        return False

    if not objects_concatable(c1.patches, c2.patches):
115
116
        return False

comfyanonymous's avatar
comfyanonymous committed
117
    return cond_equal_size(c1.conditioning, c2.conditioning)
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

def cond_cat(c_list):
    c_crossattn = []
    c_concat = []
    c_adm = []
    crossattn_max_len = 0

    temp = {}
    for x in c_list:
        for k in x:
            cur = temp.get(k, [])
            cur.append(x[k])
            temp[k] = cur

    out = {}
    for k in temp:
        conds = temp[k]
        out[k] = conds[0].concat(conds[1:])

    return out

139
140
141
def calc_cond_batch(model, conds, x_in, timestep, model_options):
    out_conds = []
    out_counts = []
142
143
    to_run = []

144
145
146
    for i in range(len(conds)):
        out_conds.append(torch.zeros_like(x_in))
        out_counts.append(torch.ones_like(x_in) * 1e-37)
147

148
149
150
151
152
153
154
155
        cond = conds[i]
        if cond is not None:
            for x in cond:
                p = get_area_and_mult(x, x_in, timestep)
                if p is None:
                    continue

                to_run += [(p, i)]
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185

    while len(to_run) > 0:
        first = to_run[0]
        first_shape = first[0][0].shape
        to_batch_temp = []
        for x in range(len(to_run)):
            if can_concat_cond(to_run[x][0], first[0]):
                to_batch_temp += [x]

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

        free_memory = model_management.get_free_memory(x_in.device)
        for i in range(1, len(to_batch_temp) + 1):
            batch_amount = to_batch_temp[:len(to_batch_temp)//i]
            input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:]
            if model.memory_required(input_shape) < free_memory:
                to_batch = batch_amount
                break

        input_x = []
        mult = []
        c = []
        cond_or_uncond = []
        area = []
        control = None
        patches = None
        for x in to_batch:
            o = to_run.pop(x)
            p = o[0]
comfyanonymous's avatar
comfyanonymous committed
186
187
188
189
190
191
192
            input_x.append(p.input_x)
            mult.append(p.mult)
            c.append(p.conditioning)
            area.append(p.area)
            cond_or_uncond.append(o[1])
            control = p.control
            patches = p.patches
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211

        batch_chunks = len(cond_or_uncond)
        input_x = torch.cat(input_x)
        c = cond_cat(c)
        timestep_ = torch.cat([timestep] * batch_chunks)

        if control is not None:
            c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))

        transformer_options = {}
        if 'transformer_options' in model_options:
            transformer_options = model_options['transformer_options'].copy()

        if patches is not None:
            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]
212
                    else:
213
                        cur_patches[p] = patches[p]
214
                transformer_options["patches"] = cur_patches
215
216
            else:
                transformer_options["patches"] = patches
217

218
219
        transformer_options["cond_or_uncond"] = cond_or_uncond[:]
        transformer_options["sigmas"] = timestep
220

221
        c['transformer_options'] = transformer_options
222

223
224
225
226
        if 'model_function_wrapper' in model_options:
            output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
        else:
            output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
comfyanonymous's avatar
comfyanonymous committed
227

228
        for o in range(batch_chunks):
229
            cond_index = cond_or_uncond[o]
230
231
232
233
234
235
236
237
238
239
240
241
242
            a = area[o]
            if a is None:
                out_conds[cond_index] += output[o] * mult[o]
                out_counts[cond_index] += mult[o]
            else:
                out_c = out_conds[cond_index]
                out_cts = out_counts[cond_index]
                dims = len(a) // 2
                for i in range(dims):
                    out_c = out_c.narrow(i + 2, a[i + dims], a[i])
                    out_cts = out_cts.narrow(i + 2, a[i + dims], a[i])
                out_c += output[o] * mult[o]
                out_cts += mult[o]
comfyanonymous's avatar
comfyanonymous committed
243

244
245
246
247
248
249
250
251
    for i in range(len(out_conds)):
        out_conds[i] /= out_counts[i]

    return out_conds

def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): #TODO: remove
    logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.")
    return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options))
comfyanonymous's avatar
comfyanonymous committed
252

253
def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
254
255
256
257
258
259
    if "sampler_cfg_function" in model_options:
        args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
        cfg_result = x - model_options["sampler_cfg_function"](args)
    else:
        cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
260

261
262
263
264
    for fn in model_options.get("sampler_post_cfg_function", []):
        args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
                "sigma": timestep, "model_options": model_options, "input": x}
        cfg_result = fn(args)
265

266
    return cfg_result
267

268
269
270
271
272
273
274
#The main sampling function shared by all the samplers
#Returns denoised
def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
    if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
        uncond_ = None
    else:
        uncond_ = uncond
275

276
277
    conds = [cond, uncond_]
    out = calc_cond_batch(model, conds, x, timestep, model_options)
278
    return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
comfyanonymous's avatar
comfyanonymous committed
279

comfyanonymous's avatar
comfyanonymous committed
280

281
class KSamplerX0Inpaint:
282
    def __init__(self, model, sigmas):
283
        self.inner_model = model
284
        self.sigmas = sigmas
285
    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
286
        if denoise_mask is not None:
287
            if "denoise_mask_function" in model_options:
288
                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
289
            latent_mask = 1. - denoise_mask
290
            x = x * denoise_mask + self.inner_model.inner_model.model_sampling.noise_scaling(sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1)), self.noise, self.latent_image) * latent_mask
291
        out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
292
        if denoise_mask is not None:
comfyanonymous's avatar
comfyanonymous committed
293
            out = out * denoise_mask + self.latent_image * latent_mask
294
        return out
295

296
297
def simple_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
298
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
299
    ss = len(s.sigmas) / steps
comfyanonymous's avatar
comfyanonymous committed
300
    for x in range(steps):
comfyanonymous's avatar
comfyanonymous committed
301
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
comfyanonymous's avatar
comfyanonymous committed
302
303
304
    sigs += [0.0]
    return torch.FloatTensor(sigs)

305
306
def ddim_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
307
    sigs = []
308
    ss = max(len(s.sigmas) // steps, 1)
comfyanonymous's avatar
comfyanonymous committed
309
310
311
312
313
    x = 1
    while x < len(s.sigmas):
        sigs += [float(s.sigmas[x])]
        x += ss
    sigs = sigs[::-1]
comfyanonymous's avatar
comfyanonymous committed
314
315
316
    sigs += [0.0]
    return torch.FloatTensor(sigs)

317
318
def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
319
320
321
322
323
324
325
326
    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)

327
328
329
    sigs = []
    for x in range(len(timesteps)):
        ts = timesteps[x]
comfyanonymous's avatar
comfyanonymous committed
330
        sigs.append(s.sigma(ts))
331
332
333
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
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

357
def resolve_areas_and_cond_masks_multidim(conditions, dims, device):
Jacob Segal's avatar
Jacob Segal committed
358
359
360
361
    # 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]
362
363
        if 'area' in c:
            area = c['area']
364
            if area[0] == "percentage":
365
                modified = c.copy()
366
367
368
369
370
371
372
373
                a = area[1:]
                a_len = len(a) // 2
                area = ()
                for d in range(len(dims)):
                    area += (max(1, round(a[d] * dims[d])),)
                for d in range(len(dims)):
                    area += (round(a[d + a_len] * dims[d]),)

374
                modified['area'] = area
375
                c = modified
376
377
                conditions[i] = c

378
379
        if 'mask' in c:
            mask = c['mask']
Jacob Segal's avatar
Jacob Segal committed
380
            mask = mask.to(device=device)
381
            modified = c.copy()
382
            if len(mask.shape) == len(dims):
Jacob Segal's avatar
Jacob Segal committed
383
                mask = mask.unsqueeze(0)
384
385
            if mask.shape[1:] != dims:
                mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=dims, mode='bilinear', align_corners=False).squeeze(1)
Jacob Segal's avatar
Jacob Segal committed
386

387
            if modified.get("set_area_to_bounds", False): #TODO: handle dim != 2
Jacob Segal's avatar
Jacob Segal committed
388
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
Jacob Segal's avatar
Jacob Segal committed
389
390
391
392
                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
393
                else:
Jacob Segal's avatar
Jacob Segal committed
394
                    box = boxes[0]
Jacob Segal's avatar
Jacob Segal committed
395
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
396
397
                    H = max(8, H)
                    W = max(8, W)
Jacob Segal's avatar
Jacob Segal committed
398
399
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area
Jacob Segal's avatar
Jacob Segal committed
400
401

            modified['mask'] = mask
402
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
403

404
405
406
407
408
def resolve_areas_and_cond_masks(conditions, h, w, device):
    logging.warning("WARNING: The comfy.samplers.resolve_areas_and_cond_masks function is deprecated please use the resolve_areas_and_cond_masks_multidim one instead.")
    return resolve_areas_and_cond_masks_multidim(conditions, [h, w], device)

def create_cond_with_same_area_if_none(conds, c): #TODO: handle dim != 2
409
    if 'area' not in c:
comfyanonymous's avatar
comfyanonymous committed
410
411
        return

412
    c_area = c['area']
comfyanonymous's avatar
comfyanonymous committed
413
414
    smallest = None
    for x in conds:
415
416
        if 'area' in x:
            a = x['area']
comfyanonymous's avatar
comfyanonymous committed
417
418
419
420
421
            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
422
                        elif 'area' not in smallest:
comfyanonymous's avatar
comfyanonymous committed
423
424
                            smallest = x
                        else:
425
                            if smallest['area'][0] * smallest['area'][1] > a[0] * a[1]:
comfyanonymous's avatar
comfyanonymous committed
426
427
428
429
430
431
                                smallest = x
        else:
            if smallest is None:
                smallest = x
    if smallest is None:
        return
432
433
    if 'area' in smallest:
        if smallest['area'] == c_area:
comfyanonymous's avatar
comfyanonymous committed
434
            return
435
436
437
438

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

440
def calculate_start_end_timesteps(model, conds):
441
    s = model.model_sampling
442
443
444
445
446
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
447
        if 'start_percent' in x:
448
            timestep_start = s.percent_to_sigma(x['start_percent'])
449
        if 'end_percent' in x:
450
            timestep_end = s.percent_to_sigma(x['end_percent'])
451
452

        if (timestep_start is not None) or (timestep_end is not None):
453
            n = x.copy()
454
455
456
457
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
458
            conds[t] = n
459

460
def pre_run_control(model, conds):
461
    s = model.model_sampling
462
463
464
465
466
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
467
        percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
468
        if 'control' in x:
469
            x['control'].pre_run(model, percent_to_timestep_function)
470

471
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
comfyanonymous's avatar
comfyanonymous committed
472
473
474
475
476
477
    cond_cnets = []
    cond_other = []
    uncond_cnets = []
    uncond_other = []
    for t in range(len(conds)):
        x = conds[t]
478
479
480
        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
481
482
483
484
            else:
                cond_other.append((x, t))
    for t in range(len(uncond)):
        x = uncond[t]
485
486
487
        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
488
489
490
491
492
493
494
495
496
            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]
497
498
        if name in o and o[name] is not None:
            n = o.copy()
499
            n[name] = uncond_fill_func(cond_cnets, x)
500
            uncond += [n]
comfyanonymous's avatar
comfyanonymous committed
501
        else:
502
            n = o.copy()
503
            n[name] = uncond_fill_func(cond_cnets, x)
504
            uncond[temp[1]] = n
505

506
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
507
508
    for t in range(len(conds)):
        x = conds[t]
509
        params = x.copy()
510
        params["device"] = device
511
        params["noise"] = noise
512
513
514
515
        default_width = None
        if len(noise.shape) >= 4: #TODO: 8 multiple should be set by the model
            default_width = noise.shape[3] * 8
        params["width"] = params.get("width", default_width)
516
517
        params["height"] = params.get("height", noise.shape[2] * 8)
        params["prompt_type"] = params.get("prompt_type", prompt_type)
518
519
520
521
522
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
523
524
525
526
527
528
        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
529
    return conds
530

comfyanonymous's avatar
comfyanonymous committed
531
532
533
534
535
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
536
537
538
        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
539

540
KSAMPLER_NAMES = ["euler", "euler_pp", "euler_ancestral", "euler_ancestral_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
comfyanonymous's avatar
comfyanonymous committed
541
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
comfyanonymous's avatar
comfyanonymous committed
542
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
comfyanonymous's avatar
comfyanonymous committed
543
                  "ipndm", "ipndm_v", "deis"]
comfyanonymous's avatar
comfyanonymous committed
544

545
546
547
548
549
class KSAMPLER(Sampler):
    def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
        self.sampler_function = sampler_function
        self.extra_options = extra_options
        self.inpaint_options = inpaint_options
comfyanonymous's avatar
comfyanonymous committed
550

551
552
    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
553
        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
554
555
556
557
558
559
        model_k.latent_image = latent_image
        if self.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
560

561
        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
562
563
564
565
566
567
568

        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)

        samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
569
        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
570
571
572
573
574
575
        return samples


def ksampler(sampler_name, extra_options={}, inpaint_options={}):
    if sampler_name == "dpm_fast":
        def dpm_fast_function(model, noise, sigmas, extra_args, callback, disable):
576
577
578
            if len(sigmas) <= 1:
                return noise

comfyanonymous's avatar
comfyanonymous committed
579
580
581
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
582
583
584
585
            total_steps = len(sigmas) - 1
            return k_diffusion_sampling.sample_dpm_fast(model, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=callback, disable=disable)
        sampler_function = dpm_fast_function
    elif sampler_name == "dpm_adaptive":
586
        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
587
588
589
            if len(sigmas) <= 1:
                return noise

590
591
592
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
593
            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options)
594
595
596
        sampler_function = dpm_adaptive_function
    else:
        sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
comfyanonymous's avatar
comfyanonymous committed
597

598
    return KSAMPLER(sampler_function, extra_options, inpaint_options)
comfyanonymous's avatar
comfyanonymous committed
599
600


601
602
603
def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None):
    for k in conds:
        conds[k] = conds[k][:]
604
        resolve_areas_and_cond_masks_multidim(conds[k], noise.shape[2:], device)
comfyanonymous's avatar
comfyanonymous committed
605

606
607
    for k in conds:
        calculate_start_end_timesteps(model, conds[k])
comfyanonymous's avatar
comfyanonymous committed
608

609
610
611
    if hasattr(model, 'extra_conds'):
        for k in conds:
            conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
comfyanonymous's avatar
comfyanonymous committed
612

613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
    #make sure each cond area has an opposite one with the same area
    for k in conds:
        for c in conds[k]:
            for kk in conds:
                if k != kk:
                    create_cond_with_same_area_if_none(conds[kk], c)

    for k in conds:
        pre_run_control(model, conds[k])

    if "positive" in conds:
        positive = conds["positive"]
        for k in conds:
            if k != "positive":
                apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), conds[k], 'control', lambda cond_cnets, x: cond_cnets[x])
                apply_empty_x_to_equal_area(positive, conds[k], 'gligen', lambda cond_cnets, x: cond_cnets[x])
comfyanonymous's avatar
comfyanonymous committed
629

630
631
    return conds

632
633
634
635
636
637
638
class CFGGuider:
    def __init__(self, model_patcher):
        self.model_patcher = model_patcher
        self.model_options = model_patcher.model_options
        self.original_conds = {}
        self.cfg = 1.0

comfyanonymous's avatar
comfyanonymous committed
639
640
    def set_conds(self, positive, negative):
        self.inner_set_conds({"positive": positive, "negative": negative})
641
642
643
644

    def set_cfg(self, cfg):
        self.cfg = cfg

comfyanonymous's avatar
comfyanonymous committed
645
646
647
648
    def inner_set_conds(self, conds):
        for k in conds:
            self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k])

649
650
651
652
653
654
655
656
657
    def __call__(self, *args, **kwargs):
        return self.predict_noise(*args, **kwargs)

    def predict_noise(self, x, timestep, model_options={}, seed=None):
        return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)

    def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
        if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
            latent_image = self.inner_model.process_latent_in(latent_image)
658

659
660
661
662
663
664
665
666
        self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)

        extra_args = {"model_options": self.model_options, "seed":seed}

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

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
comfyanonymous's avatar
comfyanonymous committed
667
668
669
        if sigmas.shape[-1] == 0:
            return latent_image

670
671
672
673
674
675
676
677
678
        self.conds = {}
        for k in self.original_conds:
            self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))

        self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds)
        device = self.model_patcher.load_device

        if denoise_mask is not None:
            denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)
679

680
681
682
        noise = noise.to(device)
        latent_image = latent_image.to(device)
        sigmas = sigmas.to(device)
683

684
        output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
comfyanonymous's avatar
comfyanonymous committed
685

686
687
688
689
690
        comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models)
        del self.inner_model
        del self.conds
        del self.loaded_models
        return output
comfyanonymous's avatar
comfyanonymous committed
691
692


693
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):
694
    cfg_guider = CFGGuider(model)
comfyanonymous's avatar
comfyanonymous committed
695
    cfg_guider.set_conds(positive, negative)
696
697
    cfg_guider.set_cfg(cfg)
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
comfyanonymous's avatar
comfyanonymous committed
698
699


comfyanonymous's avatar
comfyanonymous committed
700
701
702
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]

703
def calculate_sigmas(model_sampling, scheduler_name, steps):
comfyanonymous's avatar
comfyanonymous committed
704
    if scheduler_name == "karras":
705
        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
706
    elif scheduler_name == "exponential":
707
        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
708
    elif scheduler_name == "normal":
709
        sigmas = normal_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
710
    elif scheduler_name == "simple":
711
        sigmas = simple_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
712
    elif scheduler_name == "ddim_uniform":
713
        sigmas = ddim_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
714
    elif scheduler_name == "sgm_uniform":
715
        sigmas = normal_scheduler(model_sampling, steps, sgm=True)
comfyanonymous's avatar
comfyanonymous committed
716
    else:
717
        logging.error("error invalid scheduler {}".format(scheduler_name))
comfyanonymous's avatar
comfyanonymous committed
718
719
    return sigmas

720
def sampler_object(name):
721
    if name == "uni_pc":
comfyanonymous's avatar
comfyanonymous committed
722
        sampler = KSAMPLER(uni_pc.sample_unipc)
723
    elif name == "uni_pc_bh2":
comfyanonymous's avatar
comfyanonymous committed
724
        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
725
    elif name == "ddim":
726
        sampler = ksampler("euler", inpaint_options={"random": True})
727
728
729
730
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
731
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
732
733
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
734
    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
comfyanonymous's avatar
comfyanonymous committed
735

736
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
737
738
739
740
741
742
743
744
745
        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)
746
        self.denoise = denoise
747
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
748

comfyanonymous's avatar
comfyanonymous committed
749
750
751
752
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
753
        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
comfyanonymous's avatar
comfyanonymous committed
754
755
756
            steps += 1
            discard_penultimate_sigma = True

757
        sigmas = calculate_sigmas(self.model.get_model_object("model_sampling"), self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
758
759
760
761
762

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

comfyanonymous's avatar
comfyanonymous committed
763
764
    def set_steps(self, steps, denoise=None):
        self.steps = steps
765
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
766
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
767
        else:
comfyanonymous's avatar
comfyanonymous committed
768
769
770
771
772
773
            if denoise <= 0.0:
                self.sigmas = torch.FloatTensor([])
            else:
                new_steps = int(steps/denoise)
                sigmas = self.calculate_sigmas(new_steps).to(self.device)
                self.sigmas = sigmas[-(steps + 1):]
comfyanonymous's avatar
comfyanonymous committed
774

775
    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):
776
777
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
778

comfyanonymous's avatar
comfyanonymous committed
779
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
780
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
781
782
783
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
784
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
785
786
787
788
789
790
791
            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
792

793
        sampler = sampler_object(self.sampler)
794

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