| 1872 | """ |
| 1873 | |
| 1874 | def __init__( |
| 1875 | self, |
| 1876 | hidden_size: int, |
| 1877 | cross_attention_dim: int, |
| 1878 | rank: int = 4, |
| 1879 | attention_op: Optional[Callable] = None, |
| 1880 | network_alpha: Optional[int] = None, |
| 1881 | **kwargs, |
| 1882 | ): |
| 1883 | super().__init__() |
| 1884 | |
| 1885 | self.hidden_size = hidden_size |
| 1886 | self.cross_attention_dim = cross_attention_dim |
| 1887 | self.rank = rank |
| 1888 | self.attention_op = attention_op |
| 1889 | |
| 1890 | q_rank = kwargs.pop("q_rank", None) |
| 1891 | q_hidden_size = kwargs.pop("q_hidden_size", None) |
| 1892 | q_rank = q_rank if q_rank is not None else rank |
| 1893 | q_hidden_size = q_hidden_size if q_hidden_size is not None else hidden_size |
| 1894 | |
| 1895 | v_rank = kwargs.pop("v_rank", None) |
| 1896 | v_hidden_size = kwargs.pop("v_hidden_size", None) |
| 1897 | v_rank = v_rank if v_rank is not None else rank |
| 1898 | v_hidden_size = v_hidden_size if v_hidden_size is not None else hidden_size |
| 1899 | |
| 1900 | out_rank = kwargs.pop("out_rank", None) |
| 1901 | out_hidden_size = kwargs.pop("out_hidden_size", None) |
| 1902 | out_rank = out_rank if out_rank is not None else rank |
| 1903 | out_hidden_size = out_hidden_size if out_hidden_size is not None else hidden_size |
| 1904 | |
| 1905 | self.to_q_lora = LoRALinearLayer(q_hidden_size, q_hidden_size, q_rank, network_alpha) |
| 1906 | self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha) |
| 1907 | self.to_v_lora = LoRALinearLayer(cross_attention_dim or v_hidden_size, v_hidden_size, v_rank, network_alpha) |
| 1908 | self.to_out_lora = LoRALinearLayer(out_hidden_size, out_hidden_size, out_rank, network_alpha) |
| 1909 | |
| 1910 | def __call__(self, attn: Attention, hidden_states: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: |
| 1911 | self_cls_name = self.__class__.__name__ |