Upload BD3LM

modeling_bd3lm.py

@@ -302,7 +302,6 @@ class DDiTBlock(nn.Module):
   def __init__(self, n, block_size, dim, n_heads, cond_dim, mlp_ratio=4,
                dropout=0.1, max_seqlen=1024, attn_backend='flash_attn'):
     super().__init__()
-    self.max_seqlen = max_seqlen
     self.n = n
     self.block_size = block_size
     self.n_heads = n_heads
@@ -341,7 +340,7 @@ class DDiTBlock(nn.Module):
       qkv = self.attn_qkv(x)
     # store kv cache in a sliding window (can't exceed context len)
     if store_kv:
-      self.kv_cache = qkv[:, -(self.max_seqlen-self.block_size):]
+      self.kv_cache = qkv[:, -(self.n-self.block_size):]
 
     qkv = einops.rearrange(
       qkv,
@@ -389,8 +388,9 @@ class DDiTBlock(nn.Module):
 
     # get qkvs
     if mask is not None and not sample_mode:
-      qkv_x = self.get_qkv(x[:,:self.n], rotary_cos_sin)
-      qkv_x0 = self.get_qkv(x[:,self.n:], rotary_cos_sin)
+      n = mask.shape[-1] // 2
+      qkv_x = self.get_qkv(x[:,:n], rotary_cos_sin)
+      qkv_x0 = self.get_qkv(x[:,n:], rotary_cos_sin)
       qkv = torch.cat((qkv_x, qkv_x0), dim=1)
     else:
       qkv = self.get_qkv(x, rotary_cos_sin, store_kv=store_kv)
@@ -518,12 +518,13 @@ class DITBackbone(nn.Module):
       all_hidden_states.append(x)
     c = F.silu(self.sigma_map(sigma))
     if self.cross_attn:
-      rotary_cos_sin = self.rotary_emb(x[:, :self.n])
+      n = self.mask.shape[-1] // 2
+      rotary_cos_sin = self.rotary_emb(x[:, :n])
       mask = self.mask.to(x.device)
       # use block-causal mask only during sampling
       if sample_mode:
         mask = mask[
-          self.n:self.n+x.shape[1], self.n:self.n+x.shape[1]]
+          n:n+x.shape[1], n:n+x.shape[1]]
     else:
       mask = None
       rotary_cos_sin = self.rotary_emb(x)
@@ -540,8 +541,8 @@ class DITBackbone(nn.Module):
           all_hidden_states.append(x)
       logits = self.output_layer(x, c)
     if self.cross_attn and not sample_mode:
-      logits = logits[:, :self.n]
-      all_hidden_states = [hidden_states[:, :self.n] for hidden_states in all_hidden_states]
+      logits = logits[:, :n]
+      all_hidden_states = [hidden_states[:, :n] for hidden_states in all_hidden_states]
     return logits, all_hidden_states
 
 class BD3LM(transformers.PreTrainedModel):