| 2113 | """ |
| 2114 | |
| 2115 | def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=4, scale=1.0): |
| 2116 | super().__init__() |
| 2117 | |
| 2118 | if not hasattr(F, "scaled_dot_product_attention"): |
| 2119 | raise ImportError( |
| 2120 | f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." |
| 2121 | ) |
| 2122 | |
| 2123 | self.hidden_size = hidden_size |
| 2124 | self.cross_attention_dim = cross_attention_dim |
| 2125 | self.num_tokens = num_tokens |
| 2126 | self.scale = scale |
| 2127 | |
| 2128 | self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) |
| 2129 | self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) |
| 2130 | |
| 2131 | def __call__( |
| 2132 | self, |