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


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

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

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

51
            if 'mask' not in conds:
Jacob Segal's avatar
Jacob Segal committed
52
53
54
55
56
57
58
59
60
61
62
63
64
65
                rr = 8
                if area[2] != 0:
                    for t in range(rr):
                        mult[:,:,t:1+t,:] *= ((1.0/rr) * (t + 1))
                if (area[0] + area[2]) < x_in.shape[2]:
                    for t in range(rr):
                        mult[:,:,area[0] - 1 - t:area[0] - t,:] *= ((1.0/rr) * (t + 1))
                if area[3] != 0:
                    for t in range(rr):
                        mult[:,:,:,t:1+t] *= ((1.0/rr) * (t + 1))
                if (area[1] + area[3]) < x_in.shape[3]:
                    for t in range(rr):
                        mult[:,:,:,area[1] - 1 - t:area[1] - t] *= ((1.0/rr) * (t + 1))

comfyanonymous's avatar
comfyanonymous committed
66
            conditionning = {}
67
68
69
            model_conds = conds["model_conds"]
            for c in model_conds:
                conditionning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)
70

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

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

                patches['middle_patch'] = [gligen_patch]

            return (input_x, mult, conditionning, area, control, patches)
comfyanonymous's avatar
comfyanonymous committed
89
90

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

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

            #control
comfyanonymous's avatar
comfyanonymous committed
105
106
107
108
109
110
            if (c1[4] is None) != (c2[4] is None):
                return False
            if c1[4] is not None:
                if c1[4] is not c2[4]:
                    return False

111
112
113
114
115
116
117
            #patches
            if (c1[5] is None) != (c2[5] is None):
                return False
            if (c1[5] is not None):
                if c1[5] is not c2[5]:
                    return False

comfyanonymous's avatar
comfyanonymous committed
118
119
            return cond_equal_size(c1[2], c2[2])

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

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

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

comfyanonymous's avatar
comfyanonymous committed
138
139
            return out

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

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

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

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

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

163
                    to_run += [(p, UNCOND)]
164
165
166
167

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

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

                for i in range(1, len(to_batch_temp) + 1):
                    batch_amount = to_batch_temp[:len(to_batch_temp)//i]
                    if (len(batch_amount) * first_shape[0] * first_shape[2] * first_shape[3] < max_total_area):
                        to_batch = batch_amount
                        break
181
182
183
184
185
186

                input_x = []
                mult = []
                c = []
                cond_or_uncond = []
                area = []
comfyanonymous's avatar
comfyanonymous committed
187
                control = None
188
                patches = None
189
190
191
192
193
194
195
196
                for x in to_batch:
                    o = to_run.pop(x)
                    p = o[0]
                    input_x += [p[0]]
                    mult += [p[1]]
                    c += [p[2]]
                    area += [p[3]]
                    cond_or_uncond += [o[1]]
comfyanonymous's avatar
comfyanonymous committed
197
                    control = p[4]
198
                    patches = p[5]
199
200
201

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

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

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

                if patches is not None:
213
214
215
216
217
218
219
220
221
                    if "patches" in transformer_options:
                        cur_patches = transformer_options["patches"].copy()
                        for p in patches:
                            if p in cur_patches:
                                cur_patches[p] = cur_patches[p] + patches[p]
                            else:
                                cur_patches[p] = patches[p]
                    else:
                        transformer_options["patches"] = patches
222

223
                transformer_options["cond_or_uncond"] = cond_or_uncond[:]
224
                c['transformer_options'] = transformer_options
225

226
227
228
229
                if 'model_function_wrapper' in model_options:
                    output = model_options['model_function_wrapper'](model_function, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
                else:
                    output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
comfyanonymous's avatar
comfyanonymous committed
230
                del input_x
231
232
233
234
235
236
237
238

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

            out_cond /= out_count
            del out_count
243
244
245
246
            out_uncond /= out_uncond_count
            del out_uncond_count

            return out_cond, out_uncond
comfyanonymous's avatar
comfyanonymous committed
247
248


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

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

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

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
307
308
309
310
311
312
313
314
315
316
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)

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

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

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

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

            modified['mask'] = mask
385
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
386

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

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

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

419
420
421
422
423
424
def calculate_start_end_timesteps(model, conds):
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
425
426
427
428
        if 'start_percent' in x:
            timestep_start = model.sigma_to_t(model.t_to_sigma(torch.tensor(x['start_percent'] * 999.0)))
        if 'end_percent' in x:
            timestep_end = model.sigma_to_t(model.t_to_sigma(torch.tensor(x['end_percent'] * 999.0)))
429
430

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

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

        timestep_start = None
        timestep_end = None
        percent_to_timestep_function = lambda a: model.sigma_to_t(model.t_to_sigma(torch.tensor(a) * 999.0))
445
446
        if 'control' in x:
            x['control'].pre_run(model.inner_model.inner_model, percent_to_timestep_function)
447

448
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
comfyanonymous's avatar
comfyanonymous committed
449
450
451
452
453
454
    cond_cnets = []
    cond_other = []
    uncond_cnets = []
    uncond_other = []
    for t in range(len(conds)):
        x = conds[t]
455
456
457
        if 'area' not in x:
            if name in x and x[name] is not None:
                cond_cnets.append(x[name])
comfyanonymous's avatar
comfyanonymous committed
458
459
460
461
            else:
                cond_other.append((x, t))
    for t in range(len(uncond)):
        x = uncond[t]
462
463
464
        if 'area' not in x:
            if name in x and x[name] is not None:
                uncond_cnets.append(x[name])
comfyanonymous's avatar
comfyanonymous committed
465
466
467
468
469
470
471
472
473
            else:
                uncond_other.append((x, t))

    if len(uncond_cnets) > 0:
        return

    for x in range(len(cond_cnets)):
        temp = uncond_other[x % len(uncond_other)]
        o = temp[0]
474
475
        if name in o and o[name] is not None:
            n = o.copy()
476
            n[name] = uncond_fill_func(cond_cnets, x)
477
            uncond += [n]
comfyanonymous's avatar
comfyanonymous committed
478
        else:
479
            n = o.copy()
480
            n[name] = uncond_fill_func(cond_cnets, x)
481
            uncond[temp[1]] = n
482

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

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

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

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
510
511
512
        max_sigma = float(model_wrap.inner_model.model_sampling.sigma_max)
        sigma = float(sigmas[0])
        return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
comfyanonymous's avatar
comfyanonymous committed
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560

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

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

        max_denoise = self.max_denoise(model_wrap, sigmas)

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
561
def ksampler(sampler_name, extra_options={}):
comfyanonymous's avatar
comfyanonymous committed
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
    class KSAMPLER(Sampler):
        def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
            extra_args["denoise_mask"] = denoise_mask
            model_k = KSamplerX0Inpaint(model_wrap)
            model_k.latent_image = latent_image
            model_k.noise = noise

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
594
595
def wrap_model(model):
    model_denoise = CFGNoisePredictor(model)
comfyanonymous's avatar
comfyanonymous committed
596
    return model_denoise
comfyanonymous's avatar
comfyanonymous committed
597
598
599
600
601
602
603
604

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
605
    model_wrap = wrap_model(model)
comfyanonymous's avatar
comfyanonymous committed
606
607
608
609
610
611
612
613
614
615
616
617

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

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

    pre_run_control(model_wrap, negative + positive)

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

621
622
623
    if latent_image is not None:
        latent_image = model.process_latent_in(latent_image)

624
625
626
    if hasattr(model, 'extra_conds'):
        positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
        negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
comfyanonymous's avatar
comfyanonymous committed
627
628
629
630
631
632

    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
633
634
635
636
637
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
638
        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
639
    elif scheduler_name == "exponential":
comfyanonymous's avatar
comfyanonymous committed
640
        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
641
    elif scheduler_name == "normal":
comfyanonymous's avatar
comfyanonymous committed
642
        sigmas = normal_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
643
    elif scheduler_name == "simple":
comfyanonymous's avatar
comfyanonymous committed
644
        sigmas = simple_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
645
    elif scheduler_name == "ddim_uniform":
comfyanonymous's avatar
comfyanonymous committed
646
        sigmas = ddim_scheduler(model, steps)
comfyanonymous's avatar
comfyanonymous committed
647
    elif scheduler_name == "sgm_uniform":
comfyanonymous's avatar
comfyanonymous committed
648
        sigmas = normal_scheduler(model, steps, sgm=True)
comfyanonymous's avatar
comfyanonymous committed
649
650
651
652
    else:
        print("error invalid scheduler", self.scheduler)
    return sigmas

653
654
655
656
657
658
659
660
661
662
663
def sampler_class(name):
    if name == "uni_pc":
        sampler = UNIPC
    elif name == "uni_pc_bh2":
        sampler = UNIPCBH2
    elif name == "ddim":
        sampler = DDIM
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
664
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
665
666
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
comfyanonymous's avatar
comfyanonymous committed
667

668
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
669
670
671
672
673
674
675
676
677
        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)
678
        self.denoise = denoise
679
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
680

comfyanonymous's avatar
comfyanonymous committed
681
682
683
684
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
685
        if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
comfyanonymous's avatar
comfyanonymous committed
686
687
688
            steps += 1
            discard_penultimate_sigma = True

comfyanonymous's avatar
comfyanonymous committed
689
        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
690
691
692
693
694

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

comfyanonymous's avatar
comfyanonymous committed
695
696
    def set_steps(self, steps, denoise=None):
        self.steps = steps
697
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
698
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
699
700
        else:
            new_steps = int(steps/denoise)
comfyanonymous's avatar
comfyanonymous committed
701
            sigmas = self.calculate_sigmas(new_steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
702
703
            self.sigmas = sigmas[-(steps + 1):]

704
    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):
705
706
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
707

comfyanonymous's avatar
comfyanonymous committed
708
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
709
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
710
711
712
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
713
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
714
715
716
717
718
719
720
            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
721

722
        sampler = sampler_class(self.sampler)
723

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