代码
- 原始代码
def pre_dit(
self,
x: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
context_mask: Optional[torch.Tensor] = None,
action: Optional[torch.Tensor] = None,
fuse_vae_embedding_in_latents: bool = False,
control_camera_latents_input: Optional[torch.Tensor] = None,
) -> Dict[str, Any]:
x, timestep, context_mask = self._validate_forward_inputs(
x=x,
timestep=timestep,
context=context,
context_mask=context_mask,
action=action,
)
batch_size = x.shape[0]
patch_h = int(self.patch_size[1])
patch_w = int(self.patch_size[2])
if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0:
raise ValueError(
"Latent spatial shape must be divisible by DiT patch size, "
f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})"
)
tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w)
if self.seperated_timestep and fuse_vae_embedding_in_latents:
if not hasattr(self, "patch_size") or len(self.patch_size) < 3:
raise ValueError(f"Invalid dit.patch_size: {getattr(self, 'patch_size', None)}")
token_timesteps = torch.ones(
(batch_size, x.shape[2], tokens_per_frame),
dtype=timestep.dtype,
device=timestep.device,
) * timestep.view(batch_size, 1, 1)
token_timesteps[:, 0, :] = 0
token_timesteps = token_timesteps.reshape(batch_size, -1)
token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1))
t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim)
t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim))
else:
raise NotImplementedError("Only support seperated_timestep with fuse_vae_embedding_in_latents for now.")
t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim))
print("patchify前的x:", x.shape)
x = self.patchify(x, control_camera_latents_input=control_camera_latents_input)
print("patchify后的x:", x.shape)
f, h, w = x.shape[2:]
print("text_embedding前的维度:", context.shape)
context = self.text_embedding(context) # (B, L, dim)
print("text_embedding后的维度:", context.shape)
context_len = context.shape[1]
if self.action_conditioned and action is not None:
action_len = action.shape[1]
action_emb = self.action_embedding(action) # (B, action_len, dim)
action_pos_embed = sinusoidal_embedding_1d(self.hidden_dim,
torch.arange(action_len, device=action_emb.device)) # (action_len, dim)
action_emb = action_emb + action_pos_embed.unsqueeze(0) # (B, action_len, dim)
context = torch.cat([context, action_emb], dim=1) # (B, context_len + action_len, dim)
# new mask
num_temporal_groups = f - 1 # first latent frame do not attend to actions
if num_temporal_groups <= 0:
raise ValueError(
"Action-conditioned context mask requires at least 2 latent frames when `action` is provided."
)
assert action_emb.shape[1] % num_temporal_groups == 0, \
f"Action embedding length {action_emb.shape[1]} must be divisible by number of temporal groups {num_temporal_groups}"
# Each latent frame (from the 2nd one) attends to the corresponding group of action tokens
action_group_mask = create_group_causal_attn_mask(
num_temporal_groups=num_temporal_groups,
num_query_per_group=tokens_per_frame,
num_key_per_group=action_len // num_temporal_groups,
mode=self.action_group_causal_mask_mode,
).to(context.device) # ((f-1)*tokens_per_frame, action_len)
seq_len = f * h * w # query length
final_context_mask = torch.zeros((batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device) # (B, seq_len, L + action_len)
# all latent frames attend to text tokens
final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) # (B, seq_len, L)
# latent frames from the 2nd one attend to action tokens
final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand(batch_size, -1, -1) # (B, seq_len, action_len)
context_mask = final_context_mask
elif self.action_conditioned and action is None:
if f != 1:
raise ValueError(
"Action-conditioned model requires `action` unless running single-frame text-only mode with num_latent_frames=1."
)
context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L)
else:
context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L)
x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
freqs = torch.cat([
self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
], dim=-1).reshape(f * h * w, 1, -1).to(x_tokens.device)
return {
"tokens": x_tokens,
"freqs": freqs,
"t": t,
"t_mod": t_mod,
"context": context,
"context_mask": context_mask,
"meta": {
"grid_size": (f, h, w),
"tokens_per_frame": tokens_per_frame,
"batch_size": batch_size,
},
}
- 运行结果
--------before video_expert.pre_dit--------------
first_frame_latents: (1, 48, 1, 14, 28) torch.bfloat16 cuda:0
timestep_video: (1,) torch.bfloat16 cuda:0
context: (1, 129, 4096) torch.bfloat16 cuda:0
context_mask: (1, 129) torch.bool cuda:0
fuse_flag: True
patchify前的x: torch.Size([1, 48, 1, 14, 28])
patchify后的x: torch.Size([1, 3072, 1, 7, 14])
text_embedding前的维度: torch.Size([1, 129, 4096])
text_embedding后的维度: torch.Size([1, 129, 3072])
--------after video_expert.pre_dit--------------
[Processing] Key: tokens | Shape: torch.Size([1, 98, 3072]) | Dtype: torch.bfloat16
[Processing] Key: freqs | Shape: torch.Size([98, 1, 64]) | Dtype: torch.complex128
[Processing] Key: t | Shape: torch.Size([1, 98, 3072]) | Dtype: torch.bfloat16
[Processing] Key: t_mod | Shape: torch.Size([1, 98, 6, 3072]) | Dtype: torch.bfloat16
[Processing] Key: context | Shape: torch.Size([1, 129, 3072]) | Dtype: torch.bfloat16
[Processing] Key: context_mask | Shape: torch.Size([1, 98, 129]) | Dtype: torch.bool
[Processing] Key: meta | Type: <class 'dict'>
流程图
转载自 CSDN-专业IT技术社区
原文链接：https://blog.csdn.net/qq_41472205/article/details/163512441
