| 206 | |
| 207 | |
| 208 | class T5RelativeEmbedding(nn.Module): |
| 209 | def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): |
| 210 | super(T5RelativeEmbedding, self).__init__() |
| 211 | self.num_buckets = num_buckets |
| 212 | self.num_heads = num_heads |
| 213 | self.bidirectional = bidirectional |
| 214 | self.max_dist = max_dist |
| 215 | |
| 216 | # layers |
| 217 | self.embedding = nn.Embedding(num_buckets, num_heads) |
| 218 | |
| 219 | def forward(self, lq, lk): |
| 220 | device = self.embedding.weight.device |
| 221 | # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ |
| 222 | # torch.arange(lq).unsqueeze(1).to(device) |
| 223 | if torch.device(type="meta") != device: |
| 224 | rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ |
| 225 | torch.arange(lq, device=device).unsqueeze(1) |
| 226 | else: |
| 227 | rel_pos = torch.arange(lk).unsqueeze(0) - \ |
| 228 | torch.arange(lq).unsqueeze(1) |
| 229 | rel_pos = self._relative_position_bucket(rel_pos) |
| 230 | rel_pos_embeds = self.embedding(rel_pos) |
| 231 | rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( |
| 232 | 0) # [1, N, Lq, Lk] |
| 233 | return rel_pos_embeds.contiguous() |
| 234 | |
| 235 | def _relative_position_bucket(self, rel_pos): |
| 236 | # preprocess |
| 237 | if self.bidirectional: |
| 238 | num_buckets = self.num_buckets // 2 |
| 239 | rel_buckets = (rel_pos > 0).long() * num_buckets |
| 240 | rel_pos = torch.abs(rel_pos) |
| 241 | else: |
| 242 | num_buckets = self.num_buckets |
| 243 | rel_buckets = 0 |
| 244 | rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) |
| 245 | |
| 246 | # embeddings for small and large positions |
| 247 | max_exact = num_buckets // 2 |
| 248 | rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / |
| 249 | math.log(self.max_dist / max_exact) * |
| 250 | (num_buckets - max_exact)).long() |
| 251 | rel_pos_large = torch.min( |
| 252 | rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) |
| 253 | rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) |
| 254 | return rel_buckets |
| 255 | |
| 256 | class WanT5EncoderModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): |
| 257 | def __init__(self, |