(self, in_channels, n_heads, d_head,
global_cond_dim,
do_self_attention=True,
dropout=0.,
context_dim=None,
)
| 68 | """ |
| 69 | |
| 70 | def __init__(self, in_channels, n_heads, d_head, |
| 71 | global_cond_dim, |
| 72 | do_self_attention=True, |
| 73 | dropout=0., |
| 74 | context_dim=None, |
| 75 | ): |
| 76 | super().__init__() |
| 77 | |
| 78 | |
| 79 | self.in_channels = in_channels |
| 80 | inner_dim = n_heads * d_head |
| 81 | self.n_heads = n_heads |
| 82 | self.d_head = d_head |
| 83 | self.do_self_attention=do_self_attention |
| 84 | |
| 85 | |
| 86 | self.x_in_norm = AdaRMSNorm(in_channels, global_cond_dim) |
| 87 | |
| 88 | #x to qkv |
| 89 | if self.do_self_attention: |
| 90 | self.x_qkv_proj = apply_wd(torch.nn.Linear(in_channels, inner_dim * 3, bias=False)) |
| 91 | else: |
| 92 | self.x_q_proj = apply_wd(torch.nn.Linear(in_channels, inner_dim, bias=False)) |
| 93 | self.x_scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 94 | |
| 95 | self.x_pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 96 | |
| 97 | |
| 98 | #context to kv |
| 99 | self.cond_kv_proj = apply_wd(torch.nn.Linear(context_dim, inner_dim * 2, bias=False)) |
| 100 | self.cond_scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 101 | self.cond_pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 102 | |
| 103 | self.ff = FeedForwardBlock(in_channels, d_ff=int(in_channels*2), cond_features=global_cond_dim, dropout=dropout) |
| 104 | |
| 105 | self.dropout = nn.Dropout(dropout) |
| 106 | self.proj_out = apply_wd(zero_module(nn.Linear(in_channels, inner_dim))) |
| 107 | |
| 108 | |
| 109 | def forward(self, x, pos, global_cond, context=None, context_pos=None): |
nothing calls this directly
no test coverage detected