hugging-apps/direct-object-insertion

1from typing import Optional, Union2 3import torch4import torch.nn.functional as F5from torch import nn6 7from diffusers.models.attention_processor import Attention8 9 10class LoRALinearLayer(nn.Module):11    def __init__(12        self,13        in_features: int,14        out_features: int,15        rank: int = 4,16        network_alpha: Optional[float] = None,17        device: Optional[Union[torch.device, str]] = None,18        dtype: Optional[torch.dtype] = None,19        number=0,20        n_loras=1,21    ):22        super().__init__()23        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)24        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)25        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.26        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning27        self.network_alpha = network_alpha28        self.rank = rank29        self.out_features = out_features30        self.in_features = in_features31 32        nn.init.normal_(self.down.weight, std=1 / rank)33        nn.init.zeros_(self.up.weight)34 35        self.number = number36        self.n_loras = n_loras37 38    def forward(self, hidden_states: torch.Tensor, cond_seq_len: int = None) -> torch.Tensor:39        orig_dtype = hidden_states.dtype40        dtype = self.down.weight.dtype41 42        batch_size = hidden_states.shape[0]43        cond_size = cond_seq_len44 45        block_size = hidden_states.shape[1] - cond_size * self.n_loras46        shape = (batch_size, hidden_states.shape[1], 3072)47        mask = torch.ones(shape, device=hidden_states.device, dtype=dtype)48        mask[:, : block_size + self.number * cond_size, :] = 049        mask[:, block_size + (self.number + 1) * cond_size :, :] = 050        hidden_states = mask * hidden_states51 52        down_hidden_states = self.down(hidden_states.to(dtype))53        up_hidden_states = self.up(down_hidden_states)54 55        if self.network_alpha is not None:56            up_hidden_states *= self.network_alpha / self.rank57 58        return up_hidden_states.to(orig_dtype)59 60 61class TextLoRALinearLayer(nn.Module):62    def __init__(63        self,64        in_features: int,65        out_features: int,66        rank: int = 4,67        network_alpha: Optional[float] = None,68        device: Optional[Union[torch.device, str]] = None,69        dtype: Optional[torch.dtype] = None,70        token_length=512,71    ):72        super().__init__()73        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)74        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)75        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.76        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning77        self.network_alpha = network_alpha78        self.rank = rank79        self.out_features = out_features80        self.in_features = in_features81 82        nn.init.normal_(self.down.weight, std=1 / rank)83        nn.init.zeros_(self.up.weight)84 85        self.token_length = token_length86 87    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:88        orig_dtype = hidden_states.dtype89        dtype = self.down.weight.dtype90 91        batch_size, seq_len, feature_dim = hidden_states.shape92        if seq_len > self.token_length:93            mask = torch.ones((batch_size, seq_len, feature_dim), device=hidden_states.device, dtype=dtype)94            mask[:, self.token_length :, :] = 095            hidden_states = mask * hidden_states96 97        down_hidden_states = self.down(hidden_states.to(dtype))98        up_hidden_states = self.up(down_hidden_states)99 100        if self.network_alpha is not None:101            up_hidden_states *= self.network_alpha / self.rank102 103        return up_hidden_states.to(orig_dtype)104 105 106class MultiSingleStreamBlockLoraProcessor(nn.Module):107    def __init__(108        self,109        dim: int,110        ranks=[],111        lora_weights=[],112        network_alphas=[],113        device=None,114        dtype=None,115        n_loras=1,116        text_lora_config=None,117    ):118        super().__init__()119        self.n_loras = n_loras120        if text_lora_config is not None:121            self.text_len = text_lora_config.get("token_length", 512)122        else:123            self.text_len = 512124 125        self.q_loras = nn.ModuleList(126            [127                LoRALinearLayer(128                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras129                )130                for i in range(n_loras)131            ]132        )133        self.k_loras = nn.ModuleList(134            [135                LoRALinearLayer(136                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras137                )138                for i in range(n_loras)139            ]140        )141        self.v_loras = nn.ModuleList(142            [143                LoRALinearLayer(144                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras145                )146                for i in range(n_loras)147            ]148        )149        self.lora_weights = lora_weights150 151        if text_lora_config is not None:152            t_rank = text_lora_config.get("rank", 4)153            t_alpha = text_lora_config.get("alpha", None)154 155            self.text_q_lora = TextLoRALinearLayer(156                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len157            )158            self.text_k_lora = TextLoRALinearLayer(159                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len160            )161            self.text_v_lora = TextLoRALinearLayer(162                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len163            )164 165    def __call__(166        self,167        attn: Attention,168        hidden_states: torch.FloatTensor,169        encoder_hidden_states: torch.FloatTensor = None,170        attention_mask: Optional[torch.FloatTensor] = None,171        image_rotary_emb: Optional[torch.Tensor] = None,172        use_cond=False,173    ) -> torch.FloatTensor:174        batch_size, seq_len, _ = hidden_states.shape175 176        total_img_seq_len = seq_len - self.text_len177        assert total_img_seq_len % (1 + self.n_loras) == 0, (178            f"total_img_seq_len:{total_img_seq_len}, n_loras:{self.n_loras}, "179            f"seq_len:{seq_len}, text_len:{self.text_len}"180        )181        cond_seq_len = total_img_seq_len // (1 + self.n_loras)182 183        query = attn.to_q(hidden_states)184        key = attn.to_k(hidden_states)185        value = attn.to_v(hidden_states)186 187        for i in range(self.n_loras):188            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states, cond_seq_len=cond_seq_len)189            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states, cond_seq_len=cond_seq_len)190            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states, cond_seq_len=cond_seq_len)191 192        if getattr(self, "text_q_lora", None) is not None:193            query = query + self.text_q_lora(hidden_states)194        if getattr(self, "text_k_lora", None) is not None:195            key = key + self.text_k_lora(hidden_states)196        if getattr(self, "text_v_lora", None) is not None:197            value = value + self.text_v_lora(hidden_states)198 199        inner_dim = key.shape[-1]200        head_dim = inner_dim // attn.heads201 202        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)203        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)204        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)205 206        if attn.norm_q is not None:207            query = attn.norm_q(query)208        if attn.norm_k is not None:209            key = attn.norm_k(key)210 211        if image_rotary_emb is not None:212            from diffusers.models.embeddings import apply_rotary_emb213 214            query = apply_rotary_emb(query, image_rotary_emb)215            key = apply_rotary_emb(key, image_rotary_emb)216 217        cond_size = cond_seq_len218        block_size = hidden_states.shape[1] - cond_size * self.n_loras219 220        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)221 222        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)223        hidden_states = hidden_states.to(query.dtype)224 225        cond_hidden_states = hidden_states[:, block_size:, :]226        hidden_states = hidden_states[:, :block_size, :]227 228        return hidden_states if not use_cond else (hidden_states, cond_hidden_states)229 230 231class MultiDoubleStreamBlockLoraProcessor(nn.Module):232    def __init__(233        self,234        dim: int,235        ranks=[],236        lora_weights=[],237        network_alphas=[],238        device=None,239        dtype=None,240        n_loras=1,241        text_lora_config=None,242    ):243        super().__init__()244        self.n_loras = n_loras245        self.q_loras = nn.ModuleList(246            [247                LoRALinearLayer(248                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras249                )250                for i in range(n_loras)251            ]252        )253        self.k_loras = nn.ModuleList(254            [255                LoRALinearLayer(256                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras257                )258                for i in range(n_loras)259            ]260        )261        self.v_loras = nn.ModuleList(262            [263                LoRALinearLayer(264                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras265                )266                for i in range(n_loras)267            ]268        )269        self.proj_loras = nn.ModuleList(270            [271                LoRALinearLayer(272                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras273                )274                for i in range(n_loras)275            ]276        )277        self.lora_weights = lora_weights278        if text_lora_config is not None:279            t_rank = text_lora_config.get("rank", 4)280            t_alpha = text_lora_config.get("alpha", None)281            t_len = text_lora_config.get("token_length", 512)282 283            self.text_q_lora = TextLoRALinearLayer(284                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len285            )286            self.text_k_lora = TextLoRALinearLayer(287                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len288            )289            self.text_v_lora = TextLoRALinearLayer(290                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len291            )292            self.text_proj_lora = TextLoRALinearLayer(293                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len294            )295 296    def __call__(297        self,298        attn: Attention,299        hidden_states: torch.FloatTensor,300        encoder_hidden_states: torch.FloatTensor = None,301        attention_mask: Optional[torch.FloatTensor] = None,302        image_rotary_emb: Optional[torch.Tensor] = None,303        use_cond=False,304    ) -> torch.FloatTensor:305        batch_size, total_img_seq_len, _ = hidden_states.shape306 307        # `context` projections.308        inner_dim = 3072309        head_dim = inner_dim // attn.heads310        encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)311        encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)312        encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)313 314        if getattr(self, "text_q_lora", None) is not None:315            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj + self.text_q_lora(316                encoder_hidden_states317            )318        if getattr(self, "text_k_lora", None) is not None:319            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj + self.text_k_lora(encoder_hidden_states)320        if getattr(self, "text_v_lora", None) is not None:321            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj + self.text_v_lora(322                encoder_hidden_states323            )324 325        encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(326            batch_size, -1, attn.heads, head_dim327        ).transpose(1, 2)328        encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(329            batch_size, -1, attn.heads, head_dim330        ).transpose(1, 2)331        encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(332            batch_size, -1, attn.heads, head_dim333        ).transpose(1, 2)334 335        if attn.norm_added_q is not None:336            encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)337        if attn.norm_added_k is not None:338            encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)339 340        assert total_img_seq_len % (1 + self.n_loras) == 0, (341            f"total_img_seq_len:{total_img_seq_len}, n_loras:{self.n_loras}"342        )343        cond_seq_len = total_img_seq_len // (1 + self.n_loras)344 345        query = attn.to_q(hidden_states)346        key = attn.to_k(hidden_states)347        value = attn.to_v(hidden_states)348 349        for i in range(self.n_loras):350            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states, cond_seq_len=cond_seq_len)351            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states, cond_seq_len=cond_seq_len)352            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states, cond_seq_len=cond_seq_len)353 354        inner_dim = key.shape[-1]355        head_dim = inner_dim // attn.heads356        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)357        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)358        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)359 360        if attn.norm_q is not None:361            query = attn.norm_q(query)362        if attn.norm_k is not None:363            key = attn.norm_k(key)364 365        # attention366        query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)367        key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)368        value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)369 370        if image_rotary_emb is not None:371            from diffusers.models.embeddings import apply_rotary_emb372 373            query = apply_rotary_emb(query, image_rotary_emb)374            key = apply_rotary_emb(key, image_rotary_emb)375 376        cond_size = cond_seq_len377        block_size = hidden_states.shape[1] - cond_size * self.n_loras378 379        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)380 381        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)382        hidden_states = hidden_states.to(query.dtype)383 384        encoder_hidden_states, hidden_states = (385            hidden_states[:, : encoder_hidden_states.shape[1]],386            hidden_states[:, encoder_hidden_states.shape[1] :],387        )388 389        hidden_states_input = hidden_states390        hidden_states = attn.to_out[0](hidden_states)391        for i in range(self.n_loras):392            hidden_states = hidden_states + self.lora_weights[i] * self.proj_loras[i](393                hidden_states_input, cond_seq_len=cond_seq_len394            )395 396        hidden_states = attn.to_out[1](hidden_states)397 398        encoder_input = encoder_hidden_states399        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)400        if getattr(self, "text_proj_lora", None) is not None:401            encoder_hidden_states = encoder_hidden_states + self.text_proj_lora(encoder_input)402 403        cond_hidden_states = hidden_states[:, block_size:, :]404        hidden_states = hidden_states[:, :block_size, :]405 406        return (407            (hidden_states, encoder_hidden_states, cond_hidden_states)408            if use_cond409            else (encoder_hidden_states, hidden_states)410        )411