samplers.py 31.1 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
import scipy
import numpy
11

12
def get_area_and_mult(conds, x_in, timestep_in):
13
14
    dims = tuple(x_in.shape[2:])
    area = None
15
16
17
18
19
20
21
22
23
24
25
    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:
26
        area = list(conds['area'])
27
28
29
    if 'strength' in conds:
        strength = conds['strength']

30
31
32
33
34
35
    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])

36
37
38
39
40
41
42
    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']
43
44
45
46
47
48
49
50
        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
51
52
53
54
55
        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

56
    if 'mask' not in conds and area is not None:
57
        rr = 8
58
59
60
61
62
63
64
65
66
        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))
67
68
69
70
71
72

    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
73
    control = conds.get('control', None)
74
75
76
77
78
79
80
81
82
83
84
85
86
87

    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
88
89
    cond_obj = collections.namedtuple('cond_obj', ['input_x', 'mult', 'conditioning', 'area', 'control', 'patches'])
    return cond_obj(input_x, mult, conditioning, area, control, patches)
90
91
92
93
94
95
96
97
98
99
100
101

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
102
    if c1.input_x.shape != c2.input_x.shape:
103
104
        return False

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

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

    if not objects_concatable(c1.patches, c2.patches):
117
118
        return False

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

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

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

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

150
151
152
153
154
155
156
157
        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)]
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
186
187

    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
188
189
190
191
192
193
194
            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
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

        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]
214
                    else:
215
                        cur_patches[p] = patches[p]
216
                transformer_options["patches"] = cur_patches
217
218
            else:
                transformer_options["patches"] = patches
219

220
221
        transformer_options["cond_or_uncond"] = cond_or_uncond[:]
        transformer_options["sigmas"] = timestep
222

223
        c['transformer_options'] = transformer_options
224

225
226
227
228
        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
229

230
        for o in range(batch_chunks):
231
            cond_index = cond_or_uncond[o]
232
233
234
235
236
237
238
239
240
241
242
243
244
            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
245

246
247
248
249
250
251
252
253
    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
254

255
def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
256
257
258
259
260
261
    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
262

263
264
265
266
    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)
267

268
    return cfg_result
269

270
271
272
273
274
275
276
#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
277

278
279
    conds = [cond, uncond_]
    out = calc_cond_batch(model, conds, x, timestep, model_options)
280
281
282
283
284
285

    for fn in model_options.get("sampler_pre_cfg_function", []):
        args = {"conds":conds, "conds_out": out, "cond_scale": cond_scale, "timestep": timestep,
                "input": x, "sigma": timestep, "model": model, "model_options": model_options}
        out  = fn(args)

286
    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
287

comfyanonymous's avatar
comfyanonymous committed
288

289
class KSamplerX0Inpaint:
290
    def __init__(self, model, sigmas):
291
        self.inner_model = model
292
        self.sigmas = sigmas
293
    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
294
        if denoise_mask is not None:
295
            if "denoise_mask_function" in model_options:
296
                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
297
            latent_mask = 1. - denoise_mask
298
            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
299
        out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
300
        if denoise_mask is not None:
comfyanonymous's avatar
comfyanonymous committed
301
            out = out * denoise_mask + self.latent_image * latent_mask
302
        return out
303

304
305
def simple_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
306
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
307
    ss = len(s.sigmas) / steps
comfyanonymous's avatar
comfyanonymous committed
308
    for x in range(steps):
comfyanonymous's avatar
comfyanonymous committed
309
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
comfyanonymous's avatar
comfyanonymous committed
310
311
312
    sigs += [0.0]
    return torch.FloatTensor(sigs)

313
314
def ddim_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
315
    sigs = []
316
    ss = max(len(s.sigmas) // steps, 1)
comfyanonymous's avatar
comfyanonymous committed
317
318
319
320
321
    x = 1
    while x < len(s.sigmas):
        sigs += [float(s.sigmas[x])]
        x += ss
    sigs = sigs[::-1]
comfyanonymous's avatar
comfyanonymous committed
322
323
324
    sigs += [0.0]
    return torch.FloatTensor(sigs)

325
326
def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
327
328
329
330
331
332
333
334
    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)

335
336
337
    sigs = []
    for x in range(len(timesteps)):
        ts = timesteps[x]
comfyanonymous's avatar
comfyanonymous committed
338
        sigs.append(s.sigma(ts))
339
340
341
    sigs += [0.0]
    return torch.FloatTensor(sigs)

342
343
344
345
346
347
348
349
350
351
352
353
# Implemented based on: https://arxiv.org/abs/2407.12173
def beta_scheduler(model_sampling, steps, alpha=0.6, beta=0.6):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    for t in ts:
        sigs += [float(model_sampling.sigmas[int(t)])]
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
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

377
def resolve_areas_and_cond_masks_multidim(conditions, dims, device):
Jacob Segal's avatar
Jacob Segal committed
378
379
380
381
    # 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]
382
383
        if 'area' in c:
            area = c['area']
384
            if area[0] == "percentage":
385
                modified = c.copy()
386
387
388
389
390
391
392
393
                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]),)

394
                modified['area'] = area
395
                c = modified
396
397
                conditions[i] = c

398
399
        if 'mask' in c:
            mask = c['mask']
Jacob Segal's avatar
Jacob Segal committed
400
            mask = mask.to(device=device)
401
            modified = c.copy()
402
            if len(mask.shape) == len(dims):
Jacob Segal's avatar
Jacob Segal committed
403
                mask = mask.unsqueeze(0)
404
405
            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
406

407
            if modified.get("set_area_to_bounds", False): #TODO: handle dim != 2
Jacob Segal's avatar
Jacob Segal committed
408
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
Jacob Segal's avatar
Jacob Segal committed
409
410
411
412
                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
413
                else:
Jacob Segal's avatar
Jacob Segal committed
414
                    box = boxes[0]
Jacob Segal's avatar
Jacob Segal committed
415
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
416
417
                    H = max(8, H)
                    W = max(8, W)
Jacob Segal's avatar
Jacob Segal committed
418
419
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area
Jacob Segal's avatar
Jacob Segal committed
420
421

            modified['mask'] = mask
422
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
423

424
425
426
427
428
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
429
    if 'area' not in c:
comfyanonymous's avatar
comfyanonymous committed
430
431
        return

432
    c_area = c['area']
comfyanonymous's avatar
comfyanonymous committed
433
434
    smallest = None
    for x in conds:
435
436
        if 'area' in x:
            a = x['area']
comfyanonymous's avatar
comfyanonymous committed
437
438
439
440
441
            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
442
                        elif 'area' not in smallest:
comfyanonymous's avatar
comfyanonymous committed
443
444
                            smallest = x
                        else:
445
                            if smallest['area'][0] * smallest['area'][1] > a[0] * a[1]:
comfyanonymous's avatar
comfyanonymous committed
446
447
448
449
450
451
                                smallest = x
        else:
            if smallest is None:
                smallest = x
    if smallest is None:
        return
452
453
    if 'area' in smallest:
        if smallest['area'] == c_area:
comfyanonymous's avatar
comfyanonymous committed
454
            return
455
456
457
458

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

460
def calculate_start_end_timesteps(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
        if 'start_percent' in x:
468
            timestep_start = s.percent_to_sigma(x['start_percent'])
469
        if 'end_percent' in x:
470
            timestep_end = s.percent_to_sigma(x['end_percent'])
471
472

        if (timestep_start is not None) or (timestep_end is not None):
473
            n = x.copy()
474
475
476
477
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
478
            conds[t] = n
479

480
def pre_run_control(model, conds):
481
    s = model.model_sampling
482
483
484
485
486
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
487
        percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
488
        if 'control' in x:
489
            x['control'].pre_run(model, percent_to_timestep_function)
490

491
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
comfyanonymous's avatar
comfyanonymous committed
492
493
494
495
496
497
    cond_cnets = []
    cond_other = []
    uncond_cnets = []
    uncond_other = []
    for t in range(len(conds)):
        x = conds[t]
498
499
500
        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
501
502
503
504
            else:
                cond_other.append((x, t))
    for t in range(len(uncond)):
        x = uncond[t]
505
506
507
        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
508
509
510
511
512
513
514
515
516
            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]
517
518
        if name in o and o[name] is not None:
            n = o.copy()
519
            n[name] = uncond_fill_func(cond_cnets, x)
520
            uncond += [n]
comfyanonymous's avatar
comfyanonymous committed
521
        else:
522
            n = o.copy()
523
            n[name] = uncond_fill_func(cond_cnets, x)
524
            uncond[temp[1]] = n
525

526
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
527
528
    for t in range(len(conds)):
        x = conds[t]
529
        params = x.copy()
530
        params["device"] = device
531
        params["noise"] = noise
532
533
534
535
        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)
536
537
        params["height"] = params.get("height", noise.shape[2] * 8)
        params["prompt_type"] = params.get("prompt_type", prompt_type)
538
539
540
541
542
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
543
544
545
546
547
548
        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
549
    return conds
550

comfyanonymous's avatar
comfyanonymous committed
551
552
553
554
555
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
556
557
558
        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
559

560
KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
comfyanonymous's avatar
comfyanonymous committed
561
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
comfyanonymous's avatar
comfyanonymous committed
562
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
comfyanonymous's avatar
comfyanonymous committed
563
                  "ipndm", "ipndm_v", "deis"]
comfyanonymous's avatar
comfyanonymous committed
564

565
566
567
568
569
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
570

571
572
    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
573
        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
574
575
576
577
578
579
        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
580

581
        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
582
583
584
585
586
587
588

        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)
589
        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
590
591
592
593
594
595
        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):
596
597
598
            if len(sigmas) <= 1:
                return noise

comfyanonymous's avatar
comfyanonymous committed
599
600
601
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
602
603
604
605
            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":
606
        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
607
608
609
            if len(sigmas) <= 1:
                return noise

610
611
612
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
613
            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options)
614
615
616
        sampler_function = dpm_adaptive_function
    else:
        sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
comfyanonymous's avatar
comfyanonymous committed
617

618
    return KSAMPLER(sampler_function, extra_options, inpaint_options)
comfyanonymous's avatar
comfyanonymous committed
619
620


621
622
623
def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None):
    for k in conds:
        conds[k] = conds[k][:]
624
        resolve_areas_and_cond_masks_multidim(conds[k], noise.shape[2:], device)
comfyanonymous's avatar
comfyanonymous committed
625

626
627
    for k in conds:
        calculate_start_end_timesteps(model, conds[k])
comfyanonymous's avatar
comfyanonymous committed
628

629
630
631
    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
632

633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
    #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
649

650
651
    return conds

652
653
654
655
656
657
658
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
659
660
    def set_conds(self, positive, negative):
        self.inner_set_conds({"positive": positive, "negative": negative})
661
662
663
664

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

comfyanonymous's avatar
comfyanonymous committed
665
666
667
668
    def inner_set_conds(self, conds):
        for k in conds:
            self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k])

669
670
671
672
673
674
675
676
677
    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)
678

679
680
681
682
683
684
685
686
        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
687
688
689
        if sigmas.shape[-1] == 0:
            return latent_image

690
691
692
693
694
695
696
697
698
        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)
699

700
701
702
        noise = noise.to(device)
        latent_image = latent_image.to(device)
        sigmas = sigmas.to(device)
703

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

706
707
708
709
710
        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
711
712


713
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):
714
    cfg_guider = CFGGuider(model)
comfyanonymous's avatar
comfyanonymous committed
715
    cfg_guider.set_conds(positive, negative)
716
717
    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
718
719


720
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta"]
comfyanonymous's avatar
comfyanonymous committed
721
722
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]

723
def calculate_sigmas(model_sampling, scheduler_name, steps):
comfyanonymous's avatar
comfyanonymous committed
724
    if scheduler_name == "karras":
725
        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
726
    elif scheduler_name == "exponential":
727
        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
728
    elif scheduler_name == "normal":
729
        sigmas = normal_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
730
    elif scheduler_name == "simple":
731
        sigmas = simple_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
732
    elif scheduler_name == "ddim_uniform":
733
        sigmas = ddim_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
734
    elif scheduler_name == "sgm_uniform":
735
        sigmas = normal_scheduler(model_sampling, steps, sgm=True)
736
737
    elif scheduler_name == "beta":
        sigmas = beta_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
738
    else:
739
        logging.error("error invalid scheduler {}".format(scheduler_name))
comfyanonymous's avatar
comfyanonymous committed
740
741
    return sigmas

742
def sampler_object(name):
743
    if name == "uni_pc":
comfyanonymous's avatar
comfyanonymous committed
744
        sampler = KSAMPLER(uni_pc.sample_unipc)
745
    elif name == "uni_pc_bh2":
comfyanonymous's avatar
comfyanonymous committed
746
        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
747
    elif name == "ddim":
748
        sampler = ksampler("euler", inpaint_options={"random": True})
749
750
751
752
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
753
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
754
755
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
756
    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
comfyanonymous's avatar
comfyanonymous committed
757

758
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
759
760
761
762
763
764
765
766
767
        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)
768
        self.denoise = denoise
769
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
770

comfyanonymous's avatar
comfyanonymous committed
771
772
773
774
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
775
        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
comfyanonymous's avatar
comfyanonymous committed
776
777
778
            steps += 1
            discard_penultimate_sigma = True

779
        sigmas = calculate_sigmas(self.model.get_model_object("model_sampling"), self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
780
781
782
783
784

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

comfyanonymous's avatar
comfyanonymous committed
785
786
    def set_steps(self, steps, denoise=None):
        self.steps = steps
787
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
788
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
789
        else:
comfyanonymous's avatar
comfyanonymous committed
790
791
792
793
794
795
            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
796

797
    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):
798
799
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
800

comfyanonymous's avatar
comfyanonymous committed
801
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
802
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
803
804
805
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
806
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
807
808
809
810
811
812
813
            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
814

815
        sampler = sampler_object(self.sampler)
816

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