AI生成Minecraft(open-oasis) 代码浅读

mod

fixed_dims = [1] * len(shift.shape[1:])

#shift, scale 都是mod_dim
shift = shift.repeat(x.shape[0] // shift.shape[0], *fixed_dims)
scale = scale.repeat(x.shape[0] // scale.shape[0], *fixed_dims)

#sample dim可能有差异 3D <-> 2D 需要expand dim
while shift.dim() < x.dim():
    shift = shift.unsqueeze(-2)
    scale = scale.unsqueeze(-2)

return x * (1 + scale) + shift

repeat 将shift batch repeat x[0](train batch) 次, keeping dims after unmodified

gate(x, g)

control how much of one tensor (x) should pass through based on another tensor (g)

selectively turn “on” and “off” input tensor

gate 做 shift, scale 一样处理

return g * x

PatchEmbed

self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])

## 转化patch时一定要 kernel_size，stride ==  patch_size embed_dim会增大很多
## 这里 in_chan 3 -> patch_emb dim 768
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

Timestep emb

half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)

##[T, 1] * [1, dim//2] --> [T, dim//2]
args = t[:, None].float() * freqs[None]

## 补全 dim     1,2 half have same freqs
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)

FinalLayer

self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)

# Output size 为patch_size * patch_size * out_channels
# hidden_size 是flatten 是否为 patch_size * patch_size * in_channels 是否很大? emb_dim为768
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))

## Apply condition mod on final output layer
def forward(self, x, c):
    shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
    x = modulate(self.norm_final(x), shift, scale)
    x = self.linear(x)
	  return x

hidden dim做最后一次mod 之后proj

最终输出dim为 patch_size * patch_size * out_channels

SpatioTemporaDiTBlock

B, T, H, W, D = x.shape

# spatial block
s_shift_msa, s_scale_msa, s_gate_msa, s_shift_mlp, s_scale_mlp, s_gate_mlp = self.s_adaLN_modulation(c).chunk(6, dim=-1)
x = x + gate(self.s_attn(modulate(self.s_norm1(x), s_shift_msa, s_scale_msa)), s_gate_msa)
x = x + gate(self.s_mlp(modulate(self.s_norm2(x), s_shift_mlp, s_scale_mlp)), s_gate_mlp)

# temporal block
t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c).chunk(6, dim=-1)
x = x + gate(self.t_attn(modulate(self.t_norm1(x), t_shift_msa, t_scale_msa)), t_gate_msa)
x = x + gate(self.t_mlp(modulate(self.t_norm2(x), t_shift_mlp, t_scale_mlp)), t_gate_mlp)

没什么好说的先feature加spatial emb对图片自身全局总结一下再过mlp

再对加完temporal emb的自己做attend **注意只有q,k加了emb

**还要注意每次attend自己之前都做了modulate

c存储的信息从哪里来？为什么信息量足够支持生成6 * inner_feature_dim 大小

c = self.t_embedder(t)  # (N, D)
c = rearrange(c, "(b t) d -> b t d", t=T)

c其实就是 timestep信息因此会反复在每一个block中对x做condition

此外 c可能包含 external cond

if torch.is_tensor(external_cond):
    c += self.external_cond(external_cond)

还有这里面每层attend或mlp都有gate 每层都对新加入数据有量上控制同样从c产生

两组 shift, scale, gate

gate用来控制 residual shift, scale 近似 activation

在每一次 attend, mlp 之前

在final_layer时也是在 final mlp 之前确保每一层layer后都有

(adaptive activation function) 体现在生成mod的linear中 weight会被初始化成0

自适应学习mod

在squeeze-and-excitation中也会用到

DIT

input_h=18,
input_w=32,
patch_size=2,
in_channels=16,
hidden_size=1024,
depth=12,
num_heads=16,
mlp_ratio=4.0,
external_cond_dim=25,
max_frames=32,

self.x_embedder = PatchEmbed(input_h, input_w, patch_size, in_channels, hidden_size, flatten=False)
self.t_embedder = TimestepEmbedder(hidden_size)
frame_h, frame_w = self.x_embedder.grid_size

self.spatial_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads // 2, freqs_for="pixel", max_freq=256)
self.temporal_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads)
self.external_cond = nn.Linear(external_cond_dim, hidden_size) if external_cond_dim > 0 else nn.Identity()

x 先 patch_embed 到 emb_dim (b, t, h, w, d)

t 通过 t_embed 到condition (b, t) → (N, D) → (b, t, D)

c += external_cond

x, c 通过 SpatioTemporalDiTBlock

x_本身没有加 spatial_emb(针对patch排列位置)

只在spatial_axial_attend(s_attend)中对q位置进行标识?

Unpatchify

def unpatchify(self, x):
    """
    x: (N, H, W, patch_size**2 * C)
    imgs: (N, H, W, C)
    """
    c = self.out_channels
    p = self.x_embedder.patch_size[0]
    h = x.shape[1]
    w = x.shape[2]

    x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
    x = torch.einsum("nhwpqc->nchpwq", x)
    imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
    return imgs

((b, t) h, w, d) → ((b, t) h, w, p_s, p_s, c) → ((b, t) c, H, W)