| 780 | """ |
| 781 | |
| 782 | def __init__( |
| 783 | self, |
| 784 | train_kv: bool = True, |
| 785 | train_q_out: bool = True, |
| 786 | hidden_size: Optional[int] = None, |
| 787 | cross_attention_dim: Optional[int] = None, |
| 788 | out_bias: bool = True, |
| 789 | dropout: float = 0.0, |
| 790 | ): |
| 791 | super().__init__() |
| 792 | self.train_kv = train_kv |
| 793 | self.train_q_out = train_q_out |
| 794 | |
| 795 | self.hidden_size = hidden_size |
| 796 | self.cross_attention_dim = cross_attention_dim |
| 797 | |
| 798 | # `_custom_diffusion` id for easy serialization and loading. |
| 799 | if self.train_kv: |
| 800 | self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) |
| 801 | self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) |
| 802 | if self.train_q_out: |
| 803 | self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False) |
| 804 | self.to_out_custom_diffusion = nn.ModuleList([]) |
| 805 | self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias)) |
| 806 | self.to_out_custom_diffusion.append(nn.Dropout(dropout)) |
| 807 | |
| 808 | def __call__( |
| 809 | self, |