hugging-apps/echo-memory

1import torch2from .sd_text_encoder import CLIPEncoderLayer3 4 5class LoRALayerBlock(torch.nn.Module):6    def __init__(self, L, dim_in, dim_out):7        super().__init__()8        self.x = torch.nn.Parameter(torch.randn(1, L, dim_in))9        self.layer_norm = torch.nn.LayerNorm(dim_out)10 11    def forward(self, lora_A, lora_B):12        x = self.x @ lora_A.T @ lora_B.T13        x = self.layer_norm(x)14        return x15    16 17class LoRAEmbedder(torch.nn.Module):18    def __init__(self, lora_patterns=None, L=1, out_dim=2048):19        super().__init__()20        if lora_patterns is None:21            lora_patterns = self.default_lora_patterns()22            23        model_dict = {}24        for lora_pattern in lora_patterns:25            name, dim = lora_pattern["name"], lora_pattern["dim"]26            model_dict[name.replace(".", "___")] = LoRALayerBlock(L, dim[0], dim[1])27        self.model_dict = torch.nn.ModuleDict(model_dict)28        29        proj_dict = {}30        for lora_pattern in lora_patterns:31            layer_type, dim = lora_pattern["type"], lora_pattern["dim"]32            if layer_type not in proj_dict:33                proj_dict[layer_type.replace(".", "___")] = torch.nn.Linear(dim[1], out_dim)34        self.proj_dict = torch.nn.ModuleDict(proj_dict)35        36        self.lora_patterns = lora_patterns37        38        39    def default_lora_patterns(self):40        lora_patterns = []41        lora_dict = {42            "attn.a_to_qkv": (3072, 9216), "attn.a_to_out": (3072, 3072), "ff_a.0": (3072, 12288), "ff_a.2": (12288, 3072), "norm1_a.linear": (3072, 18432),43            "attn.b_to_qkv": (3072, 9216), "attn.b_to_out": (3072, 3072), "ff_b.0": (3072, 12288), "ff_b.2": (12288, 3072), "norm1_b.linear": (3072, 18432),44        }45        for i in range(19):46            for suffix in lora_dict:47                lora_patterns.append({48                    "name": f"blocks.{i}.{suffix}",49                    "dim": lora_dict[suffix],50                    "type": suffix,51                })52        lora_dict = {"to_qkv_mlp": (3072, 21504), "proj_out": (15360, 3072), "norm.linear": (3072, 9216)}53        for i in range(38):54            for suffix in lora_dict:55                lora_patterns.append({56                    "name": f"single_blocks.{i}.{suffix}",57                    "dim": lora_dict[suffix],58                    "type": suffix,59                })60        return lora_patterns61        62    def forward(self, lora):63        lora_emb = []64        for lora_pattern in self.lora_patterns:65            name, layer_type = lora_pattern["name"], lora_pattern["type"]66            lora_A = lora[name + ".lora_A.default.weight"]67            lora_B = lora[name + ".lora_B.default.weight"]68            lora_out = self.model_dict[name.replace(".", "___")](lora_A, lora_B)69            lora_out = self.proj_dict[layer_type.replace(".", "___")](lora_out)70            lora_emb.append(lora_out)71        lora_emb = torch.concat(lora_emb, dim=1)72        return lora_emb73    74    75class FluxLoRAEncoder(torch.nn.Module):76    def __init__(self, embed_dim=4096, encoder_intermediate_size=8192, num_encoder_layers=1, num_embeds_per_lora=16, num_special_embeds=1):77        super().__init__()78        self.num_embeds_per_lora = num_embeds_per_lora79        # embedder80        self.embedder = LoRAEmbedder(L=num_embeds_per_lora, out_dim=embed_dim)81        82        # encoders83        self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size, num_heads=32, head_dim=128) for _ in range(num_encoder_layers)])84 85        # special embedding86        self.special_embeds = torch.nn.Parameter(torch.randn(1, num_special_embeds, embed_dim))87        self.num_special_embeds = num_special_embeds88        89        # final layer90        self.final_layer_norm = torch.nn.LayerNorm(embed_dim)91        self.final_linear = torch.nn.Linear(embed_dim, embed_dim)92 93    def forward(self, lora):94        lora_embeds = self.embedder(lora)95        special_embeds = self.special_embeds.to(dtype=lora_embeds.dtype, device=lora_embeds.device)96        embeds = torch.concat([special_embeds, lora_embeds], dim=1)97        for encoder_id, encoder in enumerate(self.encoders):98            embeds = encoder(embeds)99        embeds = embeds[:, :self.num_special_embeds]100        embeds = self.final_layer_norm(embeds)101        embeds = self.final_linear(embeds)102        return embeds103    104    @staticmethod105    def state_dict_converter():106        return FluxLoRAEncoderStateDictConverter()107 108 109class FluxLoRAEncoderStateDictConverter:110    def from_civitai(self, state_dict):111        return state_dict112