samplers.py 29.6 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
16
#The main sampling function shared by all the samplers
#Returns predicted noise
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
264
265
266
267
268
269
270
271
272
273

class CompVisVDenoiser(k_diffusion_external.DiscreteVDDPMDenoiser):
    def __init__(self, model, quantize=False, device='cpu'):
        super().__init__(model, model.alphas_cumprod, quantize=quantize)

    def get_v(self, x, t, cond, **kwargs):
        return self.inner_model.apply_model(x, t, cond, **kwargs)


class CFGNoisePredictor(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
        self.alphas_cumprod = model.alphas_cumprod
274
275
    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
276
277
278
279
        return out


class KSamplerX0Inpaint(torch.nn.Module):
280
281
282
    def __init__(self, model):
        super().__init__()
        self.inner_model = model
283
    def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
284
285
        if denoise_mask is not None:
            latent_mask = 1. - denoise_mask
286
            x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
287
        out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
288
289
290
291
292
293
        if denoise_mask is not None:
            out *= denoise_mask

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

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

comfyanonymous's avatar
comfyanonymous committed
303
304
305
306
def ddim_scheduler(model, steps):
    sigs = []
    ddim_timesteps = make_ddim_timesteps(ddim_discr_method="uniform", num_ddim_timesteps=steps, num_ddpm_timesteps=model.inner_model.inner_model.num_timesteps, verbose=False)
    for x in range(len(ddim_timesteps) - 1, -1, -1):
307
308
309
310
        ts = ddim_timesteps[x]
        if ts > 999:
            ts = 999
        sigs.append(model.t_to_sigma(torch.tensor(ts)))
comfyanonymous's avatar
comfyanonymous committed
311
312
313
    sigs += [0.0]
    return torch.FloatTensor(sigs)

314
315
316
317
318
319
320
321
322
323
324
def sgm_scheduler(model, steps):
    sigs = []
    timesteps = torch.linspace(model.inner_model.inner_model.num_timesteps - 1, 0, steps + 1)[:-1].type(torch.int)
    for x in range(len(timesteps)):
        ts = timesteps[x]
        if ts > 999:
            ts = 999
        sigs.append(model.t_to_sigma(torch.tensor(ts)))
    sigs += [0.0]
    return torch.FloatTensor(sigs)

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

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

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

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

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

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

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

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

        timestep_start = None
        timestep_end = None
426
427
428
429
        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)))
430
431

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

439
440
441
442
443
444
445
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))
446
447
        if 'control' in x:
            x['control'].pre_run(model.inner_model.inner_model, percent_to_timestep_function)
448

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

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

        out = model_function(**params)
498
499
500
501
502
503
        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
504
    return conds
505

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

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
511
        return math.isclose(float(model_wrap.sigma_max), float(sigmas[0]), rel_tol=1e-05)
comfyanonymous's avatar
comfyanonymous committed
512
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

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
560
def ksampler(sampler_name, extra_options={}):
comfyanonymous's avatar
comfyanonymous committed
561
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
    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
589
                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
590
591
592
            return samples
    return KSAMPLER

comfyanonymous's avatar
comfyanonymous committed
593
594
595
596
597
598
599
def wrap_model(model):
    model_denoise = CFGNoisePredictor(model)
    if model.model_type == model_base.ModelType.V_PREDICTION:
        model_wrap = CompVisVDenoiser(model_denoise, quantize=True)
    else:
        model_wrap = k_diffusion_external.CompVisDenoiser(model_denoise, quantize=True)
    return model_wrap
comfyanonymous's avatar
comfyanonymous committed
600
601
602
603
604
605
606
607

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

    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)

621
    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
622
623
    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])

624
625
626
    if latent_image is not None:
        latent_image = model.process_latent_in(latent_image)

627
628
629
    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
630
631
632
633
634
635

    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
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
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):
    model_wrap = wrap_model(model)
    if scheduler_name == "karras":
        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_wrap.sigma_min), sigma_max=float(model_wrap.sigma_max))
    elif scheduler_name == "exponential":
        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_wrap.sigma_min), sigma_max=float(model_wrap.sigma_max))
    elif scheduler_name == "normal":
        sigmas = model_wrap.get_sigmas(steps)
    elif scheduler_name == "simple":
        sigmas = simple_scheduler(model_wrap, steps)
    elif scheduler_name == "ddim_uniform":
        sigmas = ddim_scheduler(model_wrap, steps)
    elif scheduler_name == "sgm_uniform":
        sigmas = sgm_scheduler(model_wrap, steps)
    else:
        print("error invalid scheduler", self.scheduler)
    return sigmas

657
658
659
660
661
662
663
664
665
666
667
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
668
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
669
670
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
comfyanonymous's avatar
comfyanonymous committed
671

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

comfyanonymous's avatar
comfyanonymous committed
685
686
687
688
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
689
        if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
comfyanonymous's avatar
comfyanonymous committed
690
691
692
            steps += 1
            discard_penultimate_sigma = True

comfyanonymous's avatar
comfyanonymous committed
693
        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
694
695
696
697
698

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

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

708
    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):
709
710
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
711

comfyanonymous's avatar
comfyanonymous committed
712
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
713
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
714
715
716
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
717
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
718
719
720
721
722
723
724
            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
725

726
        sampler = sampler_class(self.sampler)
727

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