| 1793 | """ |
| 1794 | |
| 1795 | def __init__( |
| 1796 | self, |
| 1797 | hidden_size: int, |
| 1798 | cross_attention_dim: Optional[int] = None, |
| 1799 | rank: int = 4, |
| 1800 | network_alpha: Optional[int] = None, |
| 1801 | **kwargs, |
| 1802 | ): |
| 1803 | super().__init__() |
| 1804 | if not hasattr(F, "scaled_dot_product_attention"): |
| 1805 | raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") |
| 1806 | |
| 1807 | self.hidden_size = hidden_size |
| 1808 | self.cross_attention_dim = cross_attention_dim |
| 1809 | self.rank = rank |
| 1810 | |
| 1811 | q_rank = kwargs.pop("q_rank", None) |
| 1812 | q_hidden_size = kwargs.pop("q_hidden_size", None) |
| 1813 | q_rank = q_rank if q_rank is not None else rank |
| 1814 | q_hidden_size = q_hidden_size if q_hidden_size is not None else hidden_size |
| 1815 | |
| 1816 | v_rank = kwargs.pop("v_rank", None) |
| 1817 | v_hidden_size = kwargs.pop("v_hidden_size", None) |
| 1818 | v_rank = v_rank if v_rank is not None else rank |
| 1819 | v_hidden_size = v_hidden_size if v_hidden_size is not None else hidden_size |
| 1820 | |
| 1821 | out_rank = kwargs.pop("out_rank", None) |
| 1822 | out_hidden_size = kwargs.pop("out_hidden_size", None) |
| 1823 | out_rank = out_rank if out_rank is not None else rank |
| 1824 | out_hidden_size = out_hidden_size if out_hidden_size is not None else hidden_size |
| 1825 | |
| 1826 | self.to_q_lora = LoRALinearLayer(q_hidden_size, q_hidden_size, q_rank, network_alpha) |
| 1827 | self.to_k_lora = LoRALinearLayer(cross_attention_dim or hidden_size, hidden_size, rank, network_alpha) |
| 1828 | self.to_v_lora = LoRALinearLayer(cross_attention_dim or v_hidden_size, v_hidden_size, v_rank, network_alpha) |
| 1829 | self.to_out_lora = LoRALinearLayer(out_hidden_size, out_hidden_size, out_rank, network_alpha) |
| 1830 | |
| 1831 | def __call__(self, attn: Attention, hidden_states: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor: |
| 1832 | self_cls_name = self.__class__.__name__ |