model_patcher.py 15.4 KB
Newer Older
1
2
3
4
5
import torch
import copy
import inspect

import comfy.utils
6
import comfy.model_management
7
8

class ModelPatcher:
9
    def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
10
11
12
13
        self.size = size
        self.model = model
        self.patches = {}
        self.backup = {}
14
15
        self.object_patches = {}
        self.object_patches_backup = {}
16
17
18
19
20
21
22
23
24
        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

25
26
        self.weight_inplace_update = weight_inplace_update

27
28
29
30
    def model_size(self):
        if self.size > 0:
            return self.size
        model_sd = self.model.state_dict()
31
        self.size = comfy.model_management.module_size(self.model)
32
        self.model_keys = set(model_sd.keys())
33
        return self.size
34
35

    def clone(self):
comfyanonymous's avatar
comfyanonymous committed
36
        n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)
37
38
39
40
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]

41
        n.object_patches = self.object_patches.copy()
42
43
44
45
46
47
48
49
50
        n.model_options = copy.deepcopy(self.model_options)
        n.model_keys = self.model_keys
        return n

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

51
52
53
    def memory_required(self, input_shape):
        return self.model.memory_required(input_shape=input_shape)

54
    def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
55
56
57
58
        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
59
60
        if disable_cfg1_optimization:
            self.model_options["disable_cfg1_optimization"] = True
61

62
    def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
63
        self.model_options["sampler_post_cfg_function"] = self.model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
64
65
        if disable_cfg1_optimization:
            self.model_options["disable_cfg1_optimization"] = True
66

67
68
69
70
71
72
73
74
75
    def set_model_unet_function_wrapper(self, unet_wrapper_function):
        self.model_options["model_function_wrapper"] = unet_wrapper_function

    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]

76
    def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):
77
78
79
80
81
        to = self.model_options["transformer_options"]
        if "patches_replace" not in to:
            to["patches_replace"] = {}
        if name not in to["patches_replace"]:
            to["patches_replace"][name] = {}
82
83
84
85
86
        if transformer_index is not None:
            block = (block_name, number, transformer_index)
        else:
            block = (block_name, number)
        to["patches_replace"][name][block] = patch
87
88
89
90
91
92
93

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

94
95
    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)
96

97
98
    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)
99
100
101
102
103
104
105

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

106
107
108
    def set_model_input_block_patch(self, patch):
        self.set_model_patch(patch, "input_block_patch")

109
110
111
    def set_model_input_block_patch_after_skip(self, patch):
        self.set_model_patch(patch, "input_block_patch_after_skip")

112
113
114
    def set_model_output_block_patch(self, patch):
        self.set_model_patch(patch, "output_block_patch")

115
116
117
    def add_object_patch(self, name, obj):
        self.object_patches[name] = obj

118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
    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)
134
135
        if "model_function_wrapper" in self.model_options:
            wrap_func = self.model_options["model_function_wrapper"]
136
            if hasattr(wrap_func, "to"):
137
                self.model_options["model_function_wrapper"] = wrap_func.to(device)
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

    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()
        for k in patches:
            if k in self.model_keys:
                p.add(k)
                current_patches = self.patches.get(k, [])
                current_patches.append((strength_patch, patches[k], strength_model))
                self.patches[k] = current_patches

        return list(p)

    def get_key_patches(self, filter_prefix=None):
comfyanonymous's avatar
comfyanonymous committed
155
        comfy.model_management.unload_model_clones(self)
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
        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

177
    def patch_model(self, device_to=None, patch_weights=True):
178
        for k in self.object_patches:
179
            old = comfy.utils.set_attr(self.model, k, self.object_patches[k])
180
181
182
            if k not in self.object_patches_backup:
                self.object_patches_backup[k] = old

183
184
185
186
187
188
        if patch_weights:
            model_sd = self.model_state_dict()
            for key in self.patches:
                if key not in model_sd:
                    print("could not patch. key doesn't exist in model:", key)
                    continue
189

190
                weight = model_sd[key]
191

192
                inplace_update = self.weight_inplace_update
193

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

197
198
199
200
201
202
203
204
                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:
205
                    comfy.utils.set_attr_param(self.model, key, out_weight)
206
                del temp_weight
207

208
209
210
            if device_to is not None:
                self.model.to(device_to)
                self.current_device = device_to
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226

        return self.model

    def calculate_weight(self, patches, weight, key):
        for p in patches:
            alpha = p[0]
            v = p[1]
            strength_model = p[2]

            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
227
228
229
230
231
232
                patch_type = "diff"
            elif len(v) == 2:
                patch_type = v[0]
                v = v[1]

            if patch_type == "diff":
233
234
235
236
237
                w1 = v[0]
                if alpha != 0.0:
                    if w1.shape != weight.shape:
                        print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
                    else:
238
                        weight += alpha * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype)
comfyanonymous's avatar
comfyanonymous committed
239
            elif patch_type == "lora": #lora/locon
240
241
                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)
242
243
244
245
                if v[2] is not None:
                    alpha *= v[2] / mat2.shape[0]
                if v[3] is not None:
                    #locon mid weights, hopefully the math is fine because I didn't properly test it
246
                    mat3 = comfy.model_management.cast_to_device(v[3], weight.device, torch.float32)
247
248
249
250
251
252
                    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:
                    weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
comfyanonymous's avatar
comfyanonymous committed
253
            elif patch_type == "lokr":
254
255
256
257
258
259
260
261
262
263
264
                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]
                dim = None

                if w1 is None:
                    dim = w1_b.shape[0]
265
266
                    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))
267
                else:
268
                    w1 = comfy.model_management.cast_to_device(w1, weight.device, torch.float32)
269
270
271
272

                if w2 is None:
                    dim = w2_b.shape[0]
                    if t2 is None:
273
274
                        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))
275
                    else:
276
277
278
279
                        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))
280
                else:
281
                    w2 = comfy.model_management.cast_to_device(w2, weight.device, torch.float32)
282
283
284
285
286
287
288
289
290
291

                if len(w2.shape) == 4:
                    w1 = w1.unsqueeze(2).unsqueeze(2)
                if v[2] is not None and dim is not None:
                    alpha *= v[2] / dim

                try:
                    weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
comfyanonymous's avatar
comfyanonymous committed
292
            elif patch_type == "loha":
293
294
295
296
297
298
299
300
301
                w1a = v[0]
                w1b = v[1]
                if v[2] is not None:
                    alpha *= v[2] / w1b.shape[0]
                w2a = v[3]
                w2b = v[4]
                if v[5] is not None: #cp decomposition
                    t1 = v[5]
                    t2 = v[6]
302
303
304
305
306
307
308
309
310
                    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))
311
                else:
312
313
314
315
                    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))
316
317
318
319
320

                try:
                    weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
comfyanonymous's avatar
comfyanonymous committed
321
322
323
324
325
326
327
328
329
330
            elif patch_type == "glora":
                if v[4] is not None:
                    alpha *= v[4] / v[0].shape[0]

                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)

                weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype)
comfyanonymous's avatar
comfyanonymous committed
331
332
            else:
                print("patch type not recognized", patch_type, key)
333
334
335
336
337
338

        return weight

    def unpatch_model(self, device_to=None):
        keys = list(self.backup.keys())

339
340
341
342
343
        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:
344
                comfy.utils.set_attr_param(self.model, k, self.backup[k])
345
346
347
348
349
350

        self.backup = {}

        if device_to is not None:
            self.model.to(device_to)
            self.current_device = device_to
351
352
353

        keys = list(self.object_patches_backup.keys())
        for k in keys:
354
            comfy.utils.set_attr(self.model, k, self.object_patches_backup[k])
355
356

        self.object_patches_backup = {}