| 1720 | """ |
| 1721 | |
| 1722 | def __init__( |
| 1723 | self, |
| 1724 | hidden_size: int, |
| 1725 | cross_attention_dim: Optional[int] = None, |
| 1726 | rank: int = 4, |
| 1727 | network_alpha: Optional[int] = None, |
| 1728 | **kwargs, |
| 1729 | ): |
| 1730 | super().__init__() |
| 1731 | |
| 1732 | self.hidden_size = hidden_size |
| 1733 | self.cross_attention_dim = cross_attention_dim |
| 1734 | self.rank = rank |
| 1735 | |
| 1736 | q_rank = kwargs.pop("q_rank", None) |
| 1737 | q_hidden_size = kwargs.pop("q_hidden_size", None) |
| 1738 | q_rank = q_rank if q_rank is not None else rank |
| 1739 | q_hidden_size = q_hidden_size if q_hidden_size is not None else hidden_size |
| 1740 | |
| 1741 | v_rank = kwargs.pop("v_rank", None) |
| 1742 | v_hidden_size = kwargs.pop("v_hidden_size", None) |
| 1743 | v_rank = v_rank if v_rank is not None else rank |
| 1744 | v_hidden_size = v_hidden_size if v_hidden_size is not None else hidden_size |
| 1745 | |
| 1746 | out_rank = kwargs.pop("out_rank", None) |
| 1747 | out_hidden_size = kwargs.pop("out_hidden_size", None) |
| 1748 | out_rank = out_rank if out_rank is not None else rank |
| 1749 | out_hidden_size = out_hidden_size if out_hidden_size is not None else hidden_size |
| 1750 | |
| 1751 | self.to_q_lora = LoRALinearLayer(q_hidden_size, q_hidden_size, q_rank, network_alpha) |
| 1752 | self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha) |
| 1753 | self.to_v_lora = LoRALinearLayer(cross_attention_dim or v_hidden_size, v_hidden_size, v_rank, network_alpha) |
| 1754 | self.to_out_lora = LoRALinearLayer(out_hidden_size, out_hidden_size, out_rank, network_alpha) |
| 1755 | |
| 1756 | def __call__(self, attn: Attention, hidden_states: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: |
| 1757 | self_cls_name = self.__class__.__name__ |