samplers.py 28.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
import logging
8
import comfy.sampler_helpers
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
61
62
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
63
    control = conds.get('control', None)
64
65
66
67
68
69
70
71
72
73
74
75
76
77

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

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

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

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

    if not objects_concatable(c1.patches, c2.patches):
107
108
        return False

comfyanonymous's avatar
comfyanonymous committed
109
    return cond_equal_size(c1.conditioning, c2.conditioning)
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

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

131
132
133
def calc_cond_batch(model, conds, x_in, timestep, model_options):
    out_conds = []
    out_counts = []
134
135
    to_run = []

136
137
138
    for i in range(len(conds)):
        out_conds.append(torch.zeros_like(x_in))
        out_counts.append(torch.ones_like(x_in) * 1e-37)
139

140
141
142
143
144
145
146
147
        cond = conds[i]
        if cond is not None:
            for x in cond:
                p = get_area_and_mult(x, x_in, timestep)
                if p is None:
                    continue

                to_run += [(p, i)]
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

    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
178
179
180
181
182
183
184
            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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203

        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]
204
                    else:
205
                        cur_patches[p] = patches[p]
206
                transformer_options["patches"] = cur_patches
207
208
            else:
                transformer_options["patches"] = patches
209

210
211
        transformer_options["cond_or_uncond"] = cond_or_uncond[:]
        transformer_options["sigmas"] = timestep
212

213
        c['transformer_options'] = transformer_options
214

215
216
217
218
        if 'model_function_wrapper' in model_options:
            output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
        else:
            output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
comfyanonymous's avatar
comfyanonymous committed
219

220
        for o in range(batch_chunks):
221
222
223
            cond_index = cond_or_uncond[o]
            out_conds[cond_index][:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
            out_counts[cond_index][:,:,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
224

225
226
227
228
229
230
231
232
    for i in range(len(out_conds)):
        out_conds[i] /= out_counts[i]

    return out_conds

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

234
235
236
237
238
239
240
def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}):
    if "sampler_cfg_function" in model_options:
        args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
        cfg_result = x - model_options["sampler_cfg_function"](args)
    else:
        cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
241

242
243
244
245
    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)
246

247
    return cfg_result
248

249
250
251
252
253
254
255
#The main sampling function shared by all the samplers
#Returns denoised
def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
    if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
        uncond_ = None
    else:
        uncond_ = uncond
256

257
258
259
    conds = [cond, uncond_]
    out = calc_cond_batch(model, conds, x, timestep, model_options)
    return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options)
comfyanonymous's avatar
comfyanonymous committed
260

comfyanonymous's avatar
comfyanonymous committed
261

262
class KSamplerX0Inpaint:
263
    def __init__(self, model, sigmas):
264
        self.inner_model = model
265
        self.sigmas = sigmas
266
    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
267
        if denoise_mask is not None:
268
            if "denoise_mask_function" in model_options:
269
                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
270
            latent_mask = 1. - denoise_mask
271
            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
272
        out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
273
        if denoise_mask is not None:
comfyanonymous's avatar
comfyanonymous committed
274
            out = out * denoise_mask + self.latent_image * latent_mask
275
        return out
276

277
278
def simple_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
279
    sigs = []
comfyanonymous's avatar
comfyanonymous committed
280
    ss = len(s.sigmas) / steps
comfyanonymous's avatar
comfyanonymous committed
281
    for x in range(steps):
comfyanonymous's avatar
comfyanonymous committed
282
        sigs += [float(s.sigmas[-(1 + int(x * ss))])]
comfyanonymous's avatar
comfyanonymous committed
283
284
285
    sigs += [0.0]
    return torch.FloatTensor(sigs)

286
287
def ddim_scheduler(model_sampling, steps):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
288
    sigs = []
289
    ss = max(len(s.sigmas) // steps, 1)
comfyanonymous's avatar
comfyanonymous committed
290
291
292
293
294
    x = 1
    while x < len(s.sigmas):
        sigs += [float(s.sigmas[x])]
        x += ss
    sigs = sigs[::-1]
comfyanonymous's avatar
comfyanonymous committed
295
296
297
    sigs += [0.0]
    return torch.FloatTensor(sigs)

298
299
def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
    s = model_sampling
comfyanonymous's avatar
comfyanonymous committed
300
301
302
303
304
305
306
307
    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)

308
309
310
    sigs = []
    for x in range(len(timesteps)):
        ts = timesteps[x]
comfyanonymous's avatar
comfyanonymous committed
311
        sigs.append(s.sigma(ts))
312
313
314
    sigs += [0.0]
    return torch.FloatTensor(sigs)

Jacob Segal's avatar
Jacob Segal committed
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
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

338
def resolve_areas_and_cond_masks(conditions, h, w, device):
Jacob Segal's avatar
Jacob Segal committed
339
340
341
342
    # 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]
343
344
        if 'area' in c:
            area = c['area']
345
            if area[0] == "percentage":
346
                modified = c.copy()
347
348
                area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
                modified['area'] = area
349
                c = modified
350
351
                conditions[i] = c

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

            modified['mask'] = mask
376
            conditions[i] = modified
Jacob Segal's avatar
Jacob Segal committed
377

comfyanonymous's avatar
comfyanonymous committed
378
def create_cond_with_same_area_if_none(conds, c):
379
    if 'area' not in c:
comfyanonymous's avatar
comfyanonymous committed
380
381
        return

382
    c_area = c['area']
comfyanonymous's avatar
comfyanonymous committed
383
384
    smallest = None
    for x in conds:
385
386
        if 'area' in x:
            a = x['area']
comfyanonymous's avatar
comfyanonymous committed
387
388
389
390
391
            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
392
                        elif 'area' not in smallest:
comfyanonymous's avatar
comfyanonymous committed
393
394
                            smallest = x
                        else:
395
                            if smallest['area'][0] * smallest['area'][1] > a[0] * a[1]:
comfyanonymous's avatar
comfyanonymous committed
396
397
398
399
400
401
                                smallest = x
        else:
            if smallest is None:
                smallest = x
    if smallest is None:
        return
402
403
    if 'area' in smallest:
        if smallest['area'] == c_area:
comfyanonymous's avatar
comfyanonymous committed
404
            return
405
406
407
408

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

410
def calculate_start_end_timesteps(model, conds):
411
    s = model.model_sampling
412
413
414
415
416
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
417
        if 'start_percent' in x:
418
            timestep_start = s.percent_to_sigma(x['start_percent'])
419
        if 'end_percent' in x:
420
            timestep_end = s.percent_to_sigma(x['end_percent'])
421
422

        if (timestep_start is not None) or (timestep_end is not None):
423
            n = x.copy()
424
425
426
427
            if (timestep_start is not None):
                n['timestep_start'] = timestep_start
            if (timestep_end is not None):
                n['timestep_end'] = timestep_end
428
            conds[t] = n
429

430
def pre_run_control(model, conds):
431
    s = model.model_sampling
432
433
434
435
436
    for t in range(len(conds)):
        x = conds[t]

        timestep_start = None
        timestep_end = None
437
        percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
438
        if 'control' in x:
439
            x['control'].pre_run(model, percent_to_timestep_function)
440

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

476
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
477
478
    for t in range(len(conds)):
        x = conds[t]
479
        params = x.copy()
480
        params["device"] = device
481
482
483
484
        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)
485
486
487
488
489
        for k in kwargs:
            if k not in params:
                params[k] = kwargs[k]

        out = model_function(**params)
490
491
492
493
494
495
        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
496
    return conds
497

comfyanonymous's avatar
comfyanonymous committed
498
499
500
501
502
class Sampler:
    def sample(self):
        pass

    def max_denoise(self, model_wrap, sigmas):
comfyanonymous's avatar
comfyanonymous committed
503
504
505
        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
506

comfyanonymous's avatar
comfyanonymous committed
507
KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
comfyanonymous's avatar
comfyanonymous committed
508
                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
comfyanonymous's avatar
comfyanonymous committed
509
                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
comfyanonymous's avatar
comfyanonymous committed
510

511
512
513
514
515
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
516

517
518
    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
519
        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
520
521
522
523
524
525
        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
526

527
        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
528
529
530
531
532
533
534

        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)
535
        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
536
537
538
539
540
541
        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
542
543
544
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
545
546
547
548
            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":
549
        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
550
551
552
            sigma_min = sigmas[-1]
            if sigma_min == 0:
                sigma_min = sigmas[-2]
553
            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options)
554
555
556
        sampler_function = dpm_adaptive_function
    else:
        sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
comfyanonymous's avatar
comfyanonymous committed
557

558
    return KSAMPLER(sampler_function, extra_options, inpaint_options)
comfyanonymous's avatar
comfyanonymous committed
559
560


561
562
563
564
def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None):
    for k in conds:
        conds[k] = conds[k][:]
        resolve_areas_and_cond_masks(conds[k], noise.shape[2], noise.shape[3], device)
comfyanonymous's avatar
comfyanonymous committed
565

566
567
    for k in conds:
        calculate_start_end_timesteps(model, conds[k])
comfyanonymous's avatar
comfyanonymous committed
568

569
570
571
    if hasattr(model, 'extra_conds'):
        for k in conds:
            conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
comfyanonymous's avatar
comfyanonymous committed
572

573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
    #make sure each cond area has an opposite one with the same area
    for k in conds:
        for c in conds[k]:
            for kk in conds:
                if k != kk:
                    create_cond_with_same_area_if_none(conds[kk], c)

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

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

590
591
    return conds

592
593
594
595
596
597
598
class CFGGuider:
    def __init__(self, model_patcher):
        self.model_patcher = model_patcher
        self.model_options = model_patcher.model_options
        self.original_conds = {}
        self.cfg = 1.0

comfyanonymous's avatar
comfyanonymous committed
599
600
    def set_conds(self, positive, negative):
        self.inner_set_conds({"positive": positive, "negative": negative})
601
602
603
604

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

comfyanonymous's avatar
comfyanonymous committed
605
606
607
608
    def inner_set_conds(self, conds):
        for k in conds:
            self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k])

609
610
611
612
613
614
615
616
617
    def __call__(self, *args, **kwargs):
        return self.predict_noise(*args, **kwargs)

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

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

619
620
621
622
623
624
625
626
        self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)

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

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

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

630
631
632
633
634
635
636
637
638
        self.conds = {}
        for k in self.original_conds:
            self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))

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

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

640
641
642
        noise = noise.to(device)
        latent_image = latent_image.to(device)
        sigmas = sigmas.to(device)
643

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

646
647
648
649
650
        comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models)
        del self.inner_model
        del self.conds
        del self.loaded_models
        return output
comfyanonymous's avatar
comfyanonymous committed
651
652


653
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):
654
    cfg_guider = CFGGuider(model)
comfyanonymous's avatar
comfyanonymous committed
655
    cfg_guider.set_conds(positive, negative)
656
657
    cfg_guider.set_cfg(cfg)
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
comfyanonymous's avatar
comfyanonymous committed
658
659


comfyanonymous's avatar
comfyanonymous committed
660
661
662
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]

663
def calculate_sigmas(model_sampling, scheduler_name, steps):
comfyanonymous's avatar
comfyanonymous committed
664
    if scheduler_name == "karras":
665
        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
666
    elif scheduler_name == "exponential":
667
        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
comfyanonymous's avatar
comfyanonymous committed
668
    elif scheduler_name == "normal":
669
        sigmas = normal_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
670
    elif scheduler_name == "simple":
671
        sigmas = simple_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
672
    elif scheduler_name == "ddim_uniform":
673
        sigmas = ddim_scheduler(model_sampling, steps)
comfyanonymous's avatar
comfyanonymous committed
674
    elif scheduler_name == "sgm_uniform":
675
        sigmas = normal_scheduler(model_sampling, steps, sgm=True)
comfyanonymous's avatar
comfyanonymous committed
676
    else:
677
        logging.error("error invalid scheduler {}".format(scheduler_name))
comfyanonymous's avatar
comfyanonymous committed
678
679
    return sigmas

680
def sampler_object(name):
681
    if name == "uni_pc":
comfyanonymous's avatar
comfyanonymous committed
682
        sampler = KSAMPLER(uni_pc.sample_unipc)
683
    elif name == "uni_pc_bh2":
comfyanonymous's avatar
comfyanonymous committed
684
        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
685
    elif name == "ddim":
686
        sampler = ksampler("euler", inpaint_options={"random": True})
687
688
689
690
    else:
        sampler = ksampler(name)
    return sampler

comfyanonymous's avatar
comfyanonymous committed
691
class KSampler:
comfyanonymous's avatar
comfyanonymous committed
692
693
    SCHEDULERS = SCHEDULER_NAMES
    SAMPLERS = SAMPLER_NAMES
694
    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
comfyanonymous's avatar
comfyanonymous committed
695

696
    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
comfyanonymous's avatar
comfyanonymous committed
697
698
699
700
701
702
703
704
705
        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)
706
        self.denoise = denoise
707
        self.model_options = model_options
comfyanonymous's avatar
comfyanonymous committed
708

comfyanonymous's avatar
comfyanonymous committed
709
710
711
712
    def calculate_sigmas(self, steps):
        sigmas = None

        discard_penultimate_sigma = False
713
        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
comfyanonymous's avatar
comfyanonymous committed
714
715
716
            steps += 1
            discard_penultimate_sigma = True

717
        sigmas = calculate_sigmas(self.model.get_model_object("model_sampling"), self.scheduler, steps)
comfyanonymous's avatar
comfyanonymous committed
718
719
720
721
722

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

comfyanonymous's avatar
comfyanonymous committed
723
724
    def set_steps(self, steps, denoise=None):
        self.steps = steps
725
        if denoise is None or denoise > 0.9999:
comfyanonymous's avatar
comfyanonymous committed
726
            self.sigmas = self.calculate_sigmas(steps).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
727
        else:
comfyanonymous's avatar
comfyanonymous committed
728
729
730
731
732
733
            if denoise <= 0.0:
                self.sigmas = torch.FloatTensor([])
            else:
                new_steps = int(steps/denoise)
                sigmas = self.calculate_sigmas(new_steps).to(self.device)
                self.sigmas = sigmas[-(steps + 1):]
comfyanonymous's avatar
comfyanonymous committed
734

735
    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):
736
737
        if sigmas is None:
            sigmas = self.sigmas
comfyanonymous's avatar
comfyanonymous committed
738

comfyanonymous's avatar
comfyanonymous committed
739
        if last_step is not None and last_step < (len(sigmas) - 1):
comfyanonymous's avatar
comfyanonymous committed
740
            sigmas = sigmas[:last_step + 1]
comfyanonymous's avatar
comfyanonymous committed
741
742
743
            if force_full_denoise:
                sigmas[-1] = 0

comfyanonymous's avatar
comfyanonymous committed
744
        if start_step is not None:
comfyanonymous's avatar
comfyanonymous committed
745
746
747
748
749
750
751
            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
752

753
        sampler = sampler_object(self.sampler)
754

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