(self, processor: "AttnProcessor")
| 155 | self.set_processor(processor) |
| 156 | |
| 157 | def set_processor(self, processor: "AttnProcessor"): |
| 158 | # if current processor is in `self._modules` and if passed `processor` is not, we need to |
| 159 | # pop `processor` from `self._modules` |
| 160 | if ( |
| 161 | hasattr(self, "processor") |
| 162 | and isinstance(self.processor, torch.nn.Module) |
| 163 | and not isinstance(processor, torch.nn.Module) |
| 164 | ): |
| 165 | logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}") |
| 166 | self._modules.pop("processor") |
| 167 | |
| 168 | self.processor = processor |
| 169 | |
| 170 | def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, return_attntion_probs=False, **cross_attention_kwargs): |
| 171 | # The `Attention` class can call different attention processors / attention functions |
no outgoing calls
no test coverage detected