model_patcher.py 23.7 KB
Newer Older
1
2
3
import torch
import copy
import inspect
4
import logging
5
import uuid
6
7

import comfy.utils
8
import comfy.model_management
9
10
from comfy.types import UnetWrapperFunction

11

comfyanonymous's avatar
comfyanonymous committed
12
13
14
15
def weight_decompose(dora_scale, weight, lora_diff, alpha, strength):
    dora_scale = comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32)
    lora_diff *= alpha
    weight_calc = weight + lora_diff.type(weight.dtype)
16
    weight_norm = (
comfyanonymous's avatar
comfyanonymous committed
17
18
        weight_calc.transpose(0, 1)
        .reshape(weight_calc.shape[1], -1)
19
        .norm(dim=1, keepdim=True)
comfyanonymous's avatar
comfyanonymous committed
20
        .reshape(weight_calc.shape[1], *[1] * (weight_calc.dim() - 1))
21
22
23
        .transpose(0, 1)
    )

comfyanonymous's avatar
comfyanonymous committed
24
25
26
27
28
29
30
31
    weight_calc *= (dora_scale / weight_norm).type(weight.dtype)
    if strength != 1.0:
        weight_calc -= weight
        weight += strength * (weight_calc)
    else:
        weight[:] = weight_calc
    return weight

32

comfyanonymous's avatar
comfyanonymous committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
    to = model_options["transformer_options"].copy()

    if "patches_replace" not in to:
        to["patches_replace"] = {}
    else:
        to["patches_replace"] = to["patches_replace"].copy()

    if name not in to["patches_replace"]:
        to["patches_replace"][name] = {}
    else:
        to["patches_replace"][name] = to["patches_replace"][name].copy()

    if transformer_index is not None:
        block = (block_name, number, transformer_index)
    else:
        block = (block_name, number)
    to["patches_replace"][name][block] = patch
    model_options["transformer_options"] = to
    return model_options
53

comfyanonymous's avatar
comfyanonymous committed
54
55
56
57
58
59
def set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=False):
    model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
    if disable_cfg1_optimization:
        model_options["disable_cfg1_optimization"] = True
    return model_options

60
61
62
63
64
65
def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_cfg1_optimization=False):
    model_options["sampler_pre_cfg_function"] = model_options.get("sampler_pre_cfg_function", []) + [pre_cfg_function]
    if disable_cfg1_optimization:
        model_options["disable_cfg1_optimization"] = True
    return model_options

66
class ModelPatcher:
67
    def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
68
69
70
71
        self.size = size
        self.model = model
        self.patches = {}
        self.backup = {}
72
73
        self.object_patches = {}
        self.object_patches_backup = {}
74
75
76
77
78
79
80
81
82
        self.model_options = {"transformer_options":{}}
        self.model_size()
        self.load_device = load_device
        self.offload_device = offload_device
        if current_device is None:
            self.current_device = self.offload_device
        else:
            self.current_device = current_device

83
        self.weight_inplace_update = weight_inplace_update
84
        self.model_lowvram = False
85
        self.lowvram_patch_counter = 0
86
        self.patches_uuid = uuid.uuid4()
87

88
89
90
    def model_size(self):
        if self.size > 0:
            return self.size
91
92
        self.size = comfy.model_management.module_size(self.model)
        return self.size
93
94

    def clone(self):
comfyanonymous's avatar
comfyanonymous committed
95
        n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)
96
97
98
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]
99
        n.patches_uuid = self.patches_uuid
100

101
        n.object_patches = self.object_patches.copy()
102
        n.model_options = copy.deepcopy(self.model_options)
103
104
        n.backup = self.backup
        n.object_patches_backup = self.object_patches_backup
105
106
107
108
109
110
111
        return n

    def is_clone(self, other):
        if hasattr(other, 'model') and self.model is other.model:
            return True
        return False

112
113
114
115
116
117
118
119
120
121
122
123
124
    def clone_has_same_weights(self, clone):
        if not self.is_clone(clone):
            return False

        if len(self.patches) == 0 and len(clone.patches) == 0:
            return True

        if self.patches_uuid == clone.patches_uuid:
            if len(self.patches) != len(clone.patches):
                logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.")
            else:
                return True

125
126
127
    def memory_required(self, input_shape):
        return self.model.memory_required(input_shape=input_shape)

128
    def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
129
130
131
132
        if len(inspect.signature(sampler_cfg_function).parameters) == 3:
            self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
        else:
            self.model_options["sampler_cfg_function"] = sampler_cfg_function
133
134
        if disable_cfg1_optimization:
            self.model_options["disable_cfg1_optimization"] = True
135

136
    def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
comfyanonymous's avatar
comfyanonymous committed
137
        self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)
138

139
140
141
    def set_model_sampler_pre_cfg_function(self, pre_cfg_function, disable_cfg1_optimization=False):
        self.model_options = set_model_options_pre_cfg_function(self.model_options, pre_cfg_function, disable_cfg1_optimization)

142
    def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
143
144
        self.model_options["model_function_wrapper"] = unet_wrapper_function

145
146
147
    def set_model_denoise_mask_function(self, denoise_mask_function):
        self.model_options["denoise_mask_function"] = denoise_mask_function

148
149
150
151
152
153
    def set_model_patch(self, patch, name):
        to = self.model_options["transformer_options"]
        if "patches" not in to:
            to["patches"] = {}
        to["patches"][name] = to["patches"].get(name, []) + [patch]

154
    def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):
comfyanonymous's avatar
comfyanonymous committed
155
        self.model_options = set_model_options_patch_replace(self.model_options, patch, name, block_name, number, transformer_index=transformer_index)
156
157
158
159
160
161
162

    def set_model_attn1_patch(self, patch):
        self.set_model_patch(patch, "attn1_patch")

    def set_model_attn2_patch(self, patch):
        self.set_model_patch(patch, "attn2_patch")

163
164
    def set_model_attn1_replace(self, patch, block_name, number, transformer_index=None):
        self.set_model_patch_replace(patch, "attn1", block_name, number, transformer_index)
165

166
167
    def set_model_attn2_replace(self, patch, block_name, number, transformer_index=None):
        self.set_model_patch_replace(patch, "attn2", block_name, number, transformer_index)
168
169
170
171
172
173
174

    def set_model_attn1_output_patch(self, patch):
        self.set_model_patch(patch, "attn1_output_patch")

    def set_model_attn2_output_patch(self, patch):
        self.set_model_patch(patch, "attn2_output_patch")

175
176
177
    def set_model_input_block_patch(self, patch):
        self.set_model_patch(patch, "input_block_patch")

178
179
180
    def set_model_input_block_patch_after_skip(self, patch):
        self.set_model_patch(patch, "input_block_patch_after_skip")

181
182
183
    def set_model_output_block_patch(self, patch):
        self.set_model_patch(patch, "output_block_patch")

184
185
186
    def add_object_patch(self, name, obj):
        self.object_patches[name] = obj

187
188
189
190
    def get_model_object(self, name):
        if name in self.object_patches:
            return self.object_patches[name]
        else:
191
192
193
194
            if name in self.object_patches_backup:
                return self.object_patches_backup[name]
            else:
                return comfy.utils.get_attr(self.model, name)
195

196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
    def model_patches_to(self, device):
        to = self.model_options["transformer_options"]
        if "patches" in to:
            patches = to["patches"]
            for name in patches:
                patch_list = patches[name]
                for i in range(len(patch_list)):
                    if hasattr(patch_list[i], "to"):
                        patch_list[i] = patch_list[i].to(device)
        if "patches_replace" in to:
            patches = to["patches_replace"]
            for name in patches:
                patch_list = patches[name]
                for k in patch_list:
                    if hasattr(patch_list[k], "to"):
                        patch_list[k] = patch_list[k].to(device)
212
213
        if "model_function_wrapper" in self.model_options:
            wrap_func = self.model_options["model_function_wrapper"]
214
            if hasattr(wrap_func, "to"):
215
                self.model_options["model_function_wrapper"] = wrap_func.to(device)
216
217
218
219
220
221
222

    def model_dtype(self):
        if hasattr(self.model, "get_dtype"):
            return self.model.get_dtype()

    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
        p = set()
comfyanonymous's avatar
comfyanonymous committed
223
        model_sd = self.model.state_dict()
224
        for k in patches:
225
            offset = None
226
            function = None
227
228
229
230
231
            if isinstance(k, str):
                key = k
            else:
                offset = k[1]
                key = k[0]
232
233
                if len(k) > 2:
                    function = k[2]
234
235

            if key in model_sd:
236
                p.add(k)
237
                current_patches = self.patches.get(key, [])
238
                current_patches.append((strength_patch, patches[k], strength_model, offset, function))
239
                self.patches[key] = current_patches
240

241
        self.patches_uuid = uuid.uuid4()
242
243
244
        return list(p)

    def get_key_patches(self, filter_prefix=None):
comfyanonymous's avatar
comfyanonymous committed
245
        comfy.model_management.unload_model_clones(self)
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
        model_sd = self.model_state_dict()
        p = {}
        for k in model_sd:
            if filter_prefix is not None:
                if not k.startswith(filter_prefix):
                    continue
            if k in self.patches:
                p[k] = [model_sd[k]] + self.patches[k]
            else:
                p[k] = (model_sd[k],)
        return p

    def model_state_dict(self, filter_prefix=None):
        sd = self.model.state_dict()
        keys = list(sd.keys())
        if filter_prefix is not None:
            for k in keys:
                if not k.startswith(filter_prefix):
                    sd.pop(k)
        return sd

267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
    def patch_weight_to_device(self, key, device_to=None):
        if key not in self.patches:
            return

        weight = comfy.utils.get_attr(self.model, key)

        inplace_update = self.weight_inplace_update

        if key not in self.backup:
            self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)

        if device_to is not None:
            temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
        else:
            temp_weight = weight.to(torch.float32, copy=True)
        out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
        if inplace_update:
            comfy.utils.copy_to_param(self.model, key, out_weight)
        else:
            comfy.utils.set_attr_param(self.model, key, out_weight)

288
    def patch_model(self, device_to=None, patch_weights=True):
289
        for k in self.object_patches:
290
            old = comfy.utils.set_attr(self.model, k, self.object_patches[k])
291
292
293
            if k not in self.object_patches_backup:
                self.object_patches_backup[k] = old

294
295
296
297
        if patch_weights:
            model_sd = self.model_state_dict()
            for key in self.patches:
                if key not in model_sd:
298
                    logging.warning("could not patch. key doesn't exist in model: {}".format(key))
299
                    continue
300

301
                self.patch_weight_to_device(key, device_to)
302

303
304
305
            if device_to is not None:
                self.model.to(device_to)
                self.current_device = device_to
306
307
308

        return self.model

309
    def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
310
311
312
313
314
315
316
317
318
319
320
        self.patch_model(device_to, patch_weights=False)

        logging.info("loading in lowvram mode {}".format(lowvram_model_memory/(1024 * 1024)))
        class LowVramPatch:
            def __init__(self, key, model_patcher):
                self.key = key
                self.model_patcher = model_patcher
            def __call__(self, weight):
                return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)

        mem_counter = 0
321
        patch_counter = 0
322
323
324
325
326
327
328
329
330
331
332
333
        for n, m in self.model.named_modules():
            lowvram_weight = False
            if hasattr(m, "comfy_cast_weights"):
                module_mem = comfy.model_management.module_size(m)
                if mem_counter + module_mem >= lowvram_model_memory:
                    lowvram_weight = True

            weight_key = "{}.weight".format(n)
            bias_key = "{}.bias".format(n)

            if lowvram_weight:
                if weight_key in self.patches:
334
335
336
337
                    if force_patch_weights:
                        self.patch_weight_to_device(weight_key)
                    else:
                        m.weight_function = LowVramPatch(weight_key, self)
338
                        patch_counter += 1
339
                if bias_key in self.patches:
340
341
342
343
                    if force_patch_weights:
                        self.patch_weight_to_device(bias_key)
                    else:
                        m.bias_function = LowVramPatch(bias_key, self)
344
                        patch_counter += 1
345
346
347
348
349
350
351
352
353

                m.prev_comfy_cast_weights = m.comfy_cast_weights
                m.comfy_cast_weights = True
            else:
                if hasattr(m, "weight"):
                    self.patch_weight_to_device(weight_key, device_to)
                    self.patch_weight_to_device(bias_key, device_to)
                    m.to(device_to)
                    mem_counter += comfy.model_management.module_size(m)
354
                    logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
355
356

        self.model_lowvram = True
357
        self.lowvram_patch_counter = patch_counter
358
359
        return self.model

360
361
    def calculate_weight(self, patches, weight, key):
        for p in patches:
comfyanonymous's avatar
comfyanonymous committed
362
            strength = p[0]
363
364
            v = p[1]
            strength_model = p[2]
365
            offset = p[3]
366
367
368
            function = p[4]
            if function is None:
                function = lambda a: a
369
370
371
372
373

            old_weight = None
            if offset is not None:
                old_weight = weight
                weight = weight.narrow(offset[0], offset[1], offset[2])
374
375
376
377
378
379
380
381

            if strength_model != 1.0:
                weight *= strength_model

            if isinstance(v, list):
                v = (self.calculate_weight(v[1:], v[0].clone(), key), )

            if len(v) == 1:
comfyanonymous's avatar
comfyanonymous committed
382
383
384
385
386
387
                patch_type = "diff"
            elif len(v) == 2:
                patch_type = v[0]
                v = v[1]

            if patch_type == "diff":
388
                w1 = v[0]
comfyanonymous's avatar
comfyanonymous committed
389
                if strength != 0.0:
390
                    if w1.shape != weight.shape:
391
                        logging.warning("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
392
                    else:
393
                        weight += function(strength * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype))
comfyanonymous's avatar
comfyanonymous committed
394
            elif patch_type == "lora": #lora/locon
395
396
                mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32)
                mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32)
397
                dora_scale = v[4]
398
                if v[2] is not None:
comfyanonymous's avatar
comfyanonymous committed
399
400
401
402
                    alpha = v[2] / mat2.shape[0]
                else:
                    alpha = 1.0

403
404
                if v[3] is not None:
                    #locon mid weights, hopefully the math is fine because I didn't properly test it
405
                    mat3 = comfy.model_management.cast_to_device(v[3], weight.device, torch.float32)
406
407
408
                    final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
                    mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
                try:
comfyanonymous's avatar
comfyanonymous committed
409
                    lora_diff = torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)).reshape(weight.shape)
410
                    if dora_scale is not None:
411
                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
comfyanonymous's avatar
comfyanonymous committed
412
                    else:
413
                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
414
                except Exception as e:
415
                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
comfyanonymous's avatar
comfyanonymous committed
416
            elif patch_type == "lokr":
417
418
419
420
421
422
423
                w1 = v[0]
                w2 = v[1]
                w1_a = v[3]
                w1_b = v[4]
                w2_a = v[5]
                w2_b = v[6]
                t2 = v[7]
424
                dora_scale = v[8]
425
426
427
428
                dim = None

                if w1 is None:
                    dim = w1_b.shape[0]
429
430
                    w1 = torch.mm(comfy.model_management.cast_to_device(w1_a, weight.device, torch.float32),
                                  comfy.model_management.cast_to_device(w1_b, weight.device, torch.float32))
431
                else:
432
                    w1 = comfy.model_management.cast_to_device(w1, weight.device, torch.float32)
433
434
435
436

                if w2 is None:
                    dim = w2_b.shape[0]
                    if t2 is None:
437
438
                        w2 = torch.mm(comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32),
                                      comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32))
439
                    else:
440
441
442
443
                        w2 = torch.einsum('i j k l, j r, i p -> p r k l',
                                          comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
                                          comfy.model_management.cast_to_device(w2_b, weight.device, torch.float32),
                                          comfy.model_management.cast_to_device(w2_a, weight.device, torch.float32))
444
                else:
445
                    w2 = comfy.model_management.cast_to_device(w2, weight.device, torch.float32)
446
447
448
449

                if len(w2.shape) == 4:
                    w1 = w1.unsqueeze(2).unsqueeze(2)
                if v[2] is not None and dim is not None:
comfyanonymous's avatar
comfyanonymous committed
450
451
452
                    alpha = v[2] / dim
                else:
                    alpha = 1.0
453
454

                try:
comfyanonymous's avatar
comfyanonymous committed
455
                    lora_diff = torch.kron(w1, w2).reshape(weight.shape)
456
                    if dora_scale is not None:
457
                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
comfyanonymous's avatar
comfyanonymous committed
458
                    else:
459
                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
460
                except Exception as e:
461
                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
comfyanonymous's avatar
comfyanonymous committed
462
            elif patch_type == "loha":
463
464
465
                w1a = v[0]
                w1b = v[1]
                if v[2] is not None:
comfyanonymous's avatar
comfyanonymous committed
466
467
468
469
                    alpha = v[2] / w1b.shape[0]
                else:
                    alpha = 1.0

470
471
                w2a = v[3]
                w2b = v[4]
472
                dora_scale = v[7]
473
474
475
                if v[5] is not None: #cp decomposition
                    t1 = v[5]
                    t2 = v[6]
476
477
478
479
480
481
482
483
484
                    m1 = torch.einsum('i j k l, j r, i p -> p r k l',
                                      comfy.model_management.cast_to_device(t1, weight.device, torch.float32),
                                      comfy.model_management.cast_to_device(w1b, weight.device, torch.float32),
                                      comfy.model_management.cast_to_device(w1a, weight.device, torch.float32))

                    m2 = torch.einsum('i j k l, j r, i p -> p r k l',
                                      comfy.model_management.cast_to_device(t2, weight.device, torch.float32),
                                      comfy.model_management.cast_to_device(w2b, weight.device, torch.float32),
                                      comfy.model_management.cast_to_device(w2a, weight.device, torch.float32))
485
                else:
486
487
488
489
                    m1 = torch.mm(comfy.model_management.cast_to_device(w1a, weight.device, torch.float32),
                                  comfy.model_management.cast_to_device(w1b, weight.device, torch.float32))
                    m2 = torch.mm(comfy.model_management.cast_to_device(w2a, weight.device, torch.float32),
                                  comfy.model_management.cast_to_device(w2b, weight.device, torch.float32))
490
491

                try:
comfyanonymous's avatar
comfyanonymous committed
492
                    lora_diff = (m1 * m2).reshape(weight.shape)
493
                    if dora_scale is not None:
494
                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
comfyanonymous's avatar
comfyanonymous committed
495
                    else:
496
                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
497
                except Exception as e:
498
                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
comfyanonymous's avatar
comfyanonymous committed
499
500
            elif patch_type == "glora":
                if v[4] is not None:
comfyanonymous's avatar
comfyanonymous committed
501
502
503
                    alpha = v[4] / v[0].shape[0]
                else:
                    alpha = 1.0
comfyanonymous's avatar
comfyanonymous committed
504

505
506
                dora_scale = v[5]

comfyanonymous's avatar
comfyanonymous committed
507
508
509
510
511
                a1 = comfy.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
                a2 = comfy.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
                b1 = comfy.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
                b2 = comfy.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)

512
                try:
comfyanonymous's avatar
comfyanonymous committed
513
                    lora_diff = (torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)).reshape(weight.shape)
514
                    if dora_scale is not None:
515
                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
comfyanonymous's avatar
comfyanonymous committed
516
                    else:
517
                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
518
519
                except Exception as e:
                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
comfyanonymous's avatar
comfyanonymous committed
520
            else:
521
                logging.warning("patch type not recognized {} {}".format(patch_type, key))
522

523
524
525
            if old_weight is not None:
                weight = old_weight

526
527
        return weight

528
529
530
531
532
533
534
535
536
    def unpatch_model(self, device_to=None, unpatch_weights=True):
        if unpatch_weights:
            if self.model_lowvram:
                for m in self.model.modules():
                    if hasattr(m, "prev_comfy_cast_weights"):
                        m.comfy_cast_weights = m.prev_comfy_cast_weights
                        del m.prev_comfy_cast_weights
                    m.weight_function = None
                    m.bias_function = None
537

538
                self.model_lowvram = False
539
                self.lowvram_patch_counter = 0
540

541
            keys = list(self.backup.keys())
542

543
544
545
546
547
548
            if self.weight_inplace_update:
                for k in keys:
                    comfy.utils.copy_to_param(self.model, k, self.backup[k])
            else:
                for k in keys:
                    comfy.utils.set_attr_param(self.model, k, self.backup[k])
549

550
            self.backup.clear()
551

552
553
554
            if device_to is not None:
                self.model.to(device_to)
                self.current_device = device_to
555
556
557

        keys = list(self.object_patches_backup.keys())
        for k in keys:
558
            comfy.utils.set_attr(self.model, k, self.object_patches_backup[k])
559

560
        self.object_patches_backup.clear()