brunvelop/ComfyUI

1import comfy.utils2import folder_paths3import torch4 5def load_hypernetwork_patch(path, strength):6    sd = comfy.utils.load_torch_file(path, safe_load=True)7    activation_func = sd.get('activation_func', 'linear')8    is_layer_norm = sd.get('is_layer_norm', False)9    use_dropout = sd.get('use_dropout', False)10    activate_output = sd.get('activate_output', False)11    last_layer_dropout = sd.get('last_layer_dropout', False)12 13    valid_activation = {14        "linear": torch.nn.Identity,15        "relu": torch.nn.ReLU,16        "leakyrelu": torch.nn.LeakyReLU,17        "elu": torch.nn.ELU,18        "swish": torch.nn.Hardswish,19        "tanh": torch.nn.Tanh,20        "sigmoid": torch.nn.Sigmoid,21        "softsign": torch.nn.Softsign,22        "mish": torch.nn.Mish,23    }24 25    if activation_func not in valid_activation:26        print("Unsupported Hypernetwork format, if you report it I might implement it.", path, " ", activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout)27        return None28 29    out = {}30 31    for d in sd:32        try:33            dim = int(d)34        except:35            continue36 37        output = []38        for index in [0, 1]:39            attn_weights = sd[dim][index]40            keys = attn_weights.keys()41 42            linears = filter(lambda a: a.endswith(".weight"), keys)43            linears = list(map(lambda a: a[:-len(".weight")], linears))44            layers = []45 46            i = 047            while i < len(linears):48                lin_name = linears[i]49                last_layer = (i == (len(linears) - 1))50                penultimate_layer = (i == (len(linears) - 2))51 52                lin_weight = attn_weights['{}.weight'.format(lin_name)]53                lin_bias = attn_weights['{}.bias'.format(lin_name)]54                layer = torch.nn.Linear(lin_weight.shape[1], lin_weight.shape[0])55                layer.load_state_dict({"weight": lin_weight, "bias": lin_bias})56                layers.append(layer)57                if activation_func != "linear":58                    if (not last_layer) or (activate_output):59                        layers.append(valid_activation[activation_func]())60                if is_layer_norm:61                    i += 162                    ln_name = linears[i]63                    ln_weight = attn_weights['{}.weight'.format(ln_name)]64                    ln_bias = attn_weights['{}.bias'.format(ln_name)]65                    ln = torch.nn.LayerNorm(ln_weight.shape[0])66                    ln.load_state_dict({"weight": ln_weight, "bias": ln_bias})67                    layers.append(ln)68                if use_dropout:69                    if (not last_layer) and (not penultimate_layer or last_layer_dropout):70                        layers.append(torch.nn.Dropout(p=0.3))71                i += 172 73            output.append(torch.nn.Sequential(*layers))74        out[dim] = torch.nn.ModuleList(output)75 76    class hypernetwork_patch:77        def __init__(self, hypernet, strength):78            self.hypernet = hypernet79            self.strength = strength80        def __call__(self, q, k, v, extra_options):81            dim = k.shape[-1]82            if dim in self.hypernet:83                hn = self.hypernet[dim]84                k = k + hn[0](k) * self.strength85                v = v + hn[1](v) * self.strength86 87            return q, k, v88 89        def to(self, device):90            for d in self.hypernet.keys():91                self.hypernet[d] = self.hypernet[d].to(device)92            return self93 94    return hypernetwork_patch(out, strength)95 96class HypernetworkLoader:97    @classmethod98    def INPUT_TYPES(s):99        return {"required": { "model": ("MODEL",),100                              "hypernetwork_name": (folder_paths.get_filename_list("hypernetworks"), ),101                              "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),102                              }}103    RETURN_TYPES = ("MODEL",)104    FUNCTION = "load_hypernetwork"105 106    CATEGORY = "loaders"107 108    def load_hypernetwork(self, model, hypernetwork_name, strength):109        hypernetwork_path = folder_paths.get_full_path("hypernetworks", hypernetwork_name)110        model_hypernetwork = model.clone()111        patch = load_hypernetwork_patch(hypernetwork_path, strength)112        if patch is not None:113            model_hypernetwork.set_model_attn1_patch(patch)114            model_hypernetwork.set_model_attn2_patch(patch)115        return (model_hypernetwork,)116 117NODE_CLASS_MAPPINGS = {118    "HypernetworkLoader": HypernetworkLoader119}120