samplers.py 26.9 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

8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def get_area_and_mult(conds, x_in, timestep_in):
    area = (x_in.shape[2], x_in.shape[3], 0, 0)
    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:
        area = conds['area']
    if 'strength' in conds:
        strength = conds['strength']

    input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
    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']
        assert(mask.shape[1] == x_in.shape[2])
        assert(mask.shape[2] == x_in.shape[3])
        mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
        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

    if 'mask' not in conds:
        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))

    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
61
    control = conds.get('control', None)
62
63
64
65
66
67
68
69
70
71
72
73
74
75

    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
76
77
    cond_obj = collections.namedtuple('cond_obj', ['input_x', 'mult', 'conditioning', 'area', 'control', 'patches'])
    return cond_obj(input_x, mult, conditioning, area, control, patches)
78
79
80
81
82
83
84
85
86
87
88
89

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
90
    if c1.input_x.shape != c2.input_x.shape:
91
92
        return False

comfyanonymous's avatar
comfyanonymous committed
93
94
    def objects_concatable(obj1, obj2):
        if (obj1 is None) != (obj2 is None):
95
            return False
comfyanonymous's avatar
comfyanonymous committed
96
97
98
99
        if obj1 is not None:
            if obj1 is not obj2:
                return False
        return True
100

comfyanonymous's avatar
comfyanonymous committed
101
102
103
104
    if not objects_concatable(c1.control, c2.control):
        return False

    if not objects_concatable(c1.patches, c2.patches):
105
106
        return False

comfyanonymous's avatar
comfyanonymous committed
107
    return cond_equal_size(c1.conditioning, c2.conditioning)
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
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

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

def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
    out_cond = torch.zeros_like(x_in)
    out_count = torch.ones_like(x_in) * 1e-37

    out_uncond = torch.zeros_like(x_in)
    out_uncond_count = torch.ones_like(x_in) * 1e-37

    COND = 0
    UNCOND = 1

    to_run = []
    for x in cond:
        p = get_area_and_mult(x, x_in, timestep)
        if p is None:
            continue

        to_run += [(p, COND)]
    if uncond is not None:
        for x in uncond:
            p = get_area_and_mult(x, x_in, timestep)
            if p is None:
                continue

            to_run += [(p, UNCOND)]

    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
183
184
185
186
187
188
189
            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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208

        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]
209
                    else:
210
                        cur_patches[p] = patches[p]
211
                transformer_options["patches"] = cur_patches
212
213
            else:
                transformer_options["patches"] = patches
214

215
216
        transformer_options["cond_or_uncond"] = cond_or_uncond[:]
        transformer_options["sigmas"] = timestep
217

218
        c['transformer_options'] = transformer_options
219

220
221
222
223
224
        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)
        del input_x
comfyanonymous's avatar
comfyanonymous committed
225

226
227
228
229
230
231
232
233
        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]
        del mult
comfyanonymous's avatar
comfyanonymous committed
234

235
236
237
238
239
    out_cond /= out_count
    del out_count
    out_uncond /= out_uncond_count
    del out_uncond_count
    return out_cond, out_uncond
comfyanonymous's avatar
comfyanonymous committed
240

241
242
243
#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):
244
        if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
245
246
247
            uncond_ = None
        else:
            uncond_ = uncond
248

249
        cond_pred, uncond_pred = calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
250
        if "sampler_cfg_function" in model_options:
Hari's avatar
Hari committed
251
252
            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}
253
            cfg_result = x - model_options["sampler_cfg_function"](args)
254
255
        else:
            cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
256

257
258
259
260
        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)
261

262
        return cfg_result
comfyanonymous's avatar
comfyanonymous committed
263

comfyanonymous's avatar
comfyanonymous committed
264
265
266
267
class CFGNoisePredictor(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
268
    def apply_model(self, x, timestep, cond, uncond, cond_scale, model_options={}, seed=None):
269
        out = sampling_function(self.inner_model, x, timestep, uncond, cond, cond_scale, model_options=model_options, seed=seed)
comfyanonymous's avatar
comfyanonymous committed
270
        return out
comfyanonymous's avatar
comfyanonymous committed
271
272
    def forward(self, *args, **kwargs):
        return self.apply_model(*args, **kwargs)
comfyanonymous's avatar
comfyanonymous committed
273
274

class KSamplerX0Inpaint(torch.nn.Module):
275
    def __init__(self, model, sigmas):
276
277
        super().__init__()
        self.inner_model = model
278
        self.sigmas = sigmas
279
    def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
280
        if denoise_mask is not None:
281
            if "denoise_mask_function" in model_options:
282
                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
283
            latent_mask = 1. - denoise_mask
284
            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
285
        out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
286
        if denoise_mask is not None:
comfyanonymous's avatar
comfyanonymous committed
287
            out = out * denoise_mask + self.latent_image * latent_mask
288
        return out
289

comfyanonymous's avatar
comfyanonymous committed
290
def simple_scheduler(model, steps):
comfyanonymous's avatar
comfyanonymous committed
291
    s = model.model_sampling
comfyanonymous's avatar
comfyanonymous committed
292
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
293
    ss = len(s.sigmas) / steps
comfyanonymous's avatar
comfyanonymous committed
294
    for x in range(steps):
comfyanonymous's avatar
comfyanonymous committed
295
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
comfyanonymous's avatar
comfyanonymous committed
296
297
298
    sigs += [0.0]
    return torch.FloatTensor(sigs)

comfyanonymous's avatar
comfyanonymous committed
299
def ddim_scheduler(model, steps):
comfyanonymous's avatar
comfyanonymous committed
300
    s = model.model_sampling
comfyanonymous's avatar
comfyanonymous committed
301
    sigs = []
302
    ss = max(len(s.sigmas) // steps, 1)
comfyanonymous's avatar
comfyanonymous committed
303
304
305
306
307
    x = 1
    while x < len(s.sigmas):
        sigs += [float(s.sigmas[x])]
        x += ss
    sigs = sigs[::-1]
comfyanonymous's avatar
comfyanonymous committed
308
309
310
    sigs += [0.0]
    return torch.FloatTensor(sigs)

comfyanonymous's avatar
comfyanonymous committed
311
312
313
314
315
316
317
318
319
320
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)

321
322
323
    sigs = []
    for x in range(len(timesteps)):
        ts = timesteps[x]
comfyanonymous's avatar
comfyanonymous committed
324
        sigs.append(s.sigma(ts))
325
326
327
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
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

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

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

            modified['mask'] = mask
389
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
390

comfyanonymous's avatar
comfyanonymous committed
391
def create_cond_with_same_area_if_none(conds, c):
392
    if 'area' not in c:
comfyanonymous's avatar
comfyanonymous committed
393
394
        return

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

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

423
def calculate_start_end_timesteps(model, conds):
424
    s = model.model_sampling
425
426
427
428
429
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
430
        if 'start_percent' in x:
431
            timestep_start = s.percent_to_sigma(x['start_percent'])
432
        if 'end_percent' in x:
433
            timestep_end = s.percent_to_sigma(x['end_percent'])
434
435

        if (timestep_start is not None) or (timestep_end is not None):
436
            n = x.copy()
437
438
439
440
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
441
            conds[t] = n
442

443
def pre_run_control(model, conds):
444
    s = model.model_sampling
445
446
447
448
449
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
450
        percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
451
        if 'control' in x:
452
            x['control'].pre_run(model, percent_to_timestep_function)
453

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

489
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
490
491
    for t in range(len(conds)):
        x = conds[t]
492
        params = x.copy()
493
        params["device"] = device
494
495
496
497
        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)
498
499
500
501
502
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
503
504
505
506
507
508
        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
509
    return conds
510

comfyanonymous's avatar
comfyanonymous committed
511
512
513
514
515
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
516
517
518
        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
519

comfyanonymous's avatar
comfyanonymous committed
520
KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
comfyanonymous's avatar
comfyanonymous committed
521
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
comfyanonymous's avatar
comfyanonymous committed
522
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
comfyanonymous's avatar
comfyanonymous committed
523

524
525
526
527
528
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
529

530
531
    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
532
        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
533
534
535
536
537
538
        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
539

540
        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
541
542
543
544
545
546
547
548
549
550
551
552
553

        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)
        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):
comfyanonymous's avatar
comfyanonymous committed
554
555
556
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
557
558
559
560
561
562
563
564
565
566
567
568
            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":
        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable):
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable)
        sampler_function = dpm_adaptive_function
    else:
        sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
comfyanonymous's avatar
comfyanonymous committed
569

570
    return KSAMPLER(sampler_function, extra_options, inpaint_options)
comfyanonymous's avatar
comfyanonymous committed
571

comfyanonymous's avatar
comfyanonymous committed
572
573
def wrap_model(model):
    model_denoise = CFGNoisePredictor(model)
comfyanonymous's avatar
comfyanonymous committed
574
    return model_denoise
comfyanonymous's avatar
comfyanonymous committed
575
576
577
578
579
580
581
582

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

585
586
    calculate_start_end_timesteps(model, negative)
    calculate_start_end_timesteps(model, positive)
comfyanonymous's avatar
comfyanonymous committed
587

588
    if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
589
590
        latent_image = model.process_latent_in(latent_image)

591
    if hasattr(model, 'extra_conds'):
592
593
        positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
        negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
594

comfyanonymous's avatar
comfyanonymous committed
595
596
597
598
599
600
    #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)

601
    pre_run_control(model, negative + positive)
comfyanonymous's avatar
comfyanonymous committed
602

603
    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
604
605
606
607
608
609
610
    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])

    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
611
612
613
614
615
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
616
        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
617
    elif scheduler_name == "exponential":
comfyanonymous's avatar
comfyanonymous committed
618
        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
619
    elif scheduler_name == "normal":
comfyanonymous's avatar
comfyanonymous committed
620
        sigmas = normal_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
621
    elif scheduler_name == "simple":
comfyanonymous's avatar
comfyanonymous committed
622
        sigmas = simple_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
623
    elif scheduler_name == "ddim_uniform":
comfyanonymous's avatar
comfyanonymous committed
624
        sigmas = ddim_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
625
    elif scheduler_name == "sgm_uniform":
comfyanonymous's avatar
comfyanonymous committed
626
        sigmas = normal_scheduler(model, steps, sgm=True)
comfyanonymous's avatar
comfyanonymous committed
627
    else:
628
        print("error invalid scheduler", scheduler_name)
comfyanonymous's avatar
comfyanonymous committed
629
630
    return sigmas

631
def sampler_object(name):
632
    if name == "uni_pc":
comfyanonymous's avatar
comfyanonymous committed
633
        sampler = KSAMPLER(uni_pc.sample_unipc)
634
    elif name == "uni_pc_bh2":
comfyanonymous's avatar
comfyanonymous committed
635
        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
636
    elif name == "ddim":
637
        sampler = ksampler("euler", inpaint_options={"random": True})
638
639
640
641
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
642
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
643
644
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
645
    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
comfyanonymous's avatar
comfyanonymous committed
646

647
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
648
649
650
651
652
653
654
655
656
        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)
657
        self.denoise = denoise
658
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
659

comfyanonymous's avatar
comfyanonymous committed
660
661
662
663
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
664
        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
comfyanonymous's avatar
comfyanonymous committed
665
666
667
            steps += 1
            discard_penultimate_sigma = True

comfyanonymous's avatar
comfyanonymous committed
668
        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
669
670
671
672
673

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

comfyanonymous's avatar
comfyanonymous committed
674
675
    def set_steps(self, steps, denoise=None):
        self.steps = steps
676
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
677
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
678
679
        else:
            new_steps = int(steps/denoise)
comfyanonymous's avatar
comfyanonymous committed
680
            sigmas = self.calculate_sigmas(new_steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
681
682
            self.sigmas = sigmas[-(steps + 1):]

683
    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):
684
685
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
686

comfyanonymous's avatar
comfyanonymous committed
687
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
688
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
689
690
691
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
692
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
693
694
695
696
697
698
699
            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
700

701
        sampler = sampler_object(self.sampler)
702

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