代码原始代码defpre_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:boolFalse,control_camera_latents_input:Optional[torch.Tensor]None,)-Dict[str,Any]:x,timestep,context_maskself._validate_forward_inputs(xx,timesteptimestep,contextcontext,context_maskcontext_mask,actionaction,)batch_sizex.shape[0]patch_hint(self.patch_size[1])patch_wint(self.patch_size[2])ifx.shape[3]%patch_h!0orx.shape[4]%patch_w!0:raiseValueError(Latent spatial shape must be divisible by DiT patch size, fgot 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)ifself.seperated_timestepandfuse_vae_embedding_in_latents:ifnothasattr(self,patch_size)orlen(self.patch_size)3:raiseValueError(fInvalid dit.patch_size:{getattr(self,patch_size,None)})token_timestepstorch.ones((batch_size,x.shape[2],tokens_per_frame),dtypetimestep.dtype,devicetimestep.device,)*timestep.view(batch_size,1,1)token_timesteps[:,0,:]0token_timestepstoken_timesteps.reshape(batch_size,-1)token_t_embsinusoidal_embedding_1d(self.freq_dim,token_timesteps.reshape(-1))tself.time_embedding(token_t_emb).reshape(batch_size,-1,self.hidden_dim)t_modself.time_projection(t).unflatten(2,(6,self.hidden_dim))else:raiseNotImplementedError(Only support seperated_timestep with fuse_vae_embedding_in_latents for now.)tself.time_embedding(sinusoidal_embedding_1d(self.freq_dim,timestep))t_modself.time_projection(t).unflatten(1,(6,self.hidden_dim))print(patchify前的x:,x.shape)xself.patchify(x,control_camera_latents_inputcontrol_camera_latents_input)print(patchify后的x:,x.shape)f,h,wx.shape[2:]print(text_embedding前的维度:,context.shape)contextself.text_embedding(context)# (B, L, dim)print(text_embedding后的维度:,context.shape)context_lencontext.shape[1]ifself.action_conditionedandactionisnotNone:action_lenaction.shape[1]action_embself.action_embedding(action)# (B, action_len, dim)action_pos_embedsinusoidal_embedding_1d(self.hidden_dim,torch.arange(action_len,deviceaction_emb.device))# (action_len, dim)action_embaction_embaction_pos_embed.unsqueeze(0)# (B, action_len, dim)contexttorch.cat([context,action_emb],dim1)# (B, context_len action_len, dim)# new masknum_temporal_groupsf-1# first latent frame do not attend to actionsifnum_temporal_groups0:raiseValueError(Action-conditioned context mask requires at least 2 latent frames when action is provided.)assertaction_emb.shape[1]%num_temporal_groups0,\fAction 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 tokensaction_group_maskcreate_group_causal_attn_mask(num_temporal_groupsnum_temporal_groups,num_query_per_grouptokens_per_frame,num_key_per_groupaction_len//num_temporal_groups,modeself.action_group_causal_mask_mode,).to(context.device)# ((f-1)*tokens_per_frame, action_len)seq_lenf*h*w# query lengthfinal_context_masktorch.zeros((batch_size,seq_len,context.shape[1]),dtypetorch.bool,devicecontext.device)# (B, seq_len, L action_len)# all latent frames attend to text tokensfinal_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 tokensfinal_context_mask[:,tokens_per_frame:,context_len:]action_group_mask.unsqueeze(0).expand(batch_size,-1,-1)# (B, seq_len, action_len)context_maskfinal_context_maskelifself.action_conditionedandactionisNone:iff!1:raiseValueError(Action-conditioned model requires action unless running single-frame text-only mode with num_latent_frames1.)context_maskcontext_mask.unsqueeze(1).expand(-1,f*h*w,-1)# (B, seq_len, L)else:context_maskcontext_mask.unsqueeze(1).expand(-1,f*h*w,-1)# (B, seq_len, L)x_tokensrearrange(x,b c f h w - b (f h w) c).contiguous()freqstorch.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:classdict流程图pre_dit inputsx (first_frame_latents): (1,48,1,14,28) bf16 cudatimestep: (1,) bf16 cudacontext: (1,129,4096) bf16 cudacontext_mask: (1,129) bool cudafuse_vae_embedding_in_latents: TrueValidate shapes_validate_forward_inputsCompute tokens_per_framepatch_size(1,2,2)H14,W28patchify grid: h7,w14tokens_per_frame 7*14 98Separated timestep (per-token)seperated_timestepTrue fuseTrueBuild token_timesteps:shape (B,F,tokens_per_frame)(1,1,98)set frame0 timestep0flatten - (1,98)sinusoidal_embedding - token_t_emb(time_embedding)- t: (1,98,3072)(time_projection)- t_mod: (1,98,6,3072)Patchify video latentsBefore: x (1,48,1,14,28)After: x (1,3072,1,7,14)Flatten to token sequencex_tokens rearrange(b,c,f,h,w)-(b,f*h*w,c)tokens: (1,98,3072)Text embeddingtext_embedding: 4096 - 3072context: (1,129,4096) - (1,129,3072)Expand context_mask to per-querycontext_mask: (1,129)- (1,seq_len,129)seq_len f*h*w 98context_mask: (1,98,129)Build rotary freqsfreqs: (seq_len,1,64) (98,1,64) complex128pre_dit outputs dictReturn payload(tokens, freqs, t, t_mod,context, context_mask,meta{grid_size(1,7,14), tokens_per_frame98, batch_size1})